def build_new_name_node(*,
old_node,
attach: bool,
new_name: str,
old_name: str = None):
# build new node from new_name
if '.' in new_name:
children = []
for part in dotted_parts(new_name):
if part == '.':
children.append(Dot())
else:
children.append(Name(part))
else:
children = [Name(new_name)]
# attach to the new node subimports from the old module
if attach and type(old_node) is Node:
original_name_size = len(dotted_parts(old_name))
for part in old_node.children[original_name_size:]:
if part.value == '.':
children.append(Dot())
else:
children.append(Name(part.value))
return Node(
type=syms.dotted_name,
children=children,
prefix=old_node.prefix,
)
def rename_transform(node: LN, capture: Capture, filename: Filename) -> None:
log.debug(f"{filename} [{list(capture)}]: {node}")
# If two keys reference the same underlying object, do not modify it twice
visited: List[LN] = []
for _key, value in capture.items():
log.debug(f"{_key}: {value}")
if value in visited:
continue
visited.append(value)
if isinstance(value, Leaf) and value.type == TOKEN.NAME:
if value.value == old_name and value.parent is not None:
value.replace(Name(new_name, prefix=value.prefix))
break
elif isinstance(value, Node):
if type_repr(value.type) == "dotted_name":
dp_old = dotted_parts(old_name)
dp_new = dotted_parts(new_name)
parts = zip(dp_old, dp_new, value.children)
for old, new, leaf in parts:
if old != leaf.value:
break
if old != new:
leaf.replace(Name(new, prefix=leaf.prefix))
if len(dp_new) < len(dp_old):
# if new path is shorter, remove excess children
del value.children[len(dp_new) : len(dp_old)]
elif len(dp_new) > len(dp_old):
# if new path is longer, add new children
children = [
Name(new) for new in dp_new[len(dp_old) : len(dp_new)]
]
value.children[len(dp_old) : len(dp_old)] = children
elif type_repr(value.type) == "power":
# We don't actually need the '.' so just skip it
dp_old = old_name.split(".")
dp_new = new_name.split(".")
for old, new, leaf in zip(dp_old, dp_new, value.children):
if isinstance(leaf, Node):
name_leaf = leaf.children[1]
else:
name_leaf = leaf
if old != name_leaf.value:
break
name_leaf.replace(Name(new, prefix=name_leaf.prefix))
if len(dp_new) < len(dp_old):
# if new path is shorter, remove excess children
del value.children[len(dp_new) : len(dp_old)]
elif len(dp_new) > len(dp_old):
# if new path is longer, add new trailers in the middle
for i in range(len(dp_old), len(dp_new)):
value.insert_child(
i, Node(SYMBOL.trailer, [Dot(), Name(dp_new[i])])
)
def handle_assertraises(node, capture, arguments):
"""
with self.assertRaises(x):
--> with pytest.raises(x):
self.assertRaises(ValueError, func, arg1)
--> pytest.raises(ValueError, func, arg1)
"""
capture['attr1'].replace(kw('pytest', prefix=capture['attr1'].prefix))
capture['attr2'].replace(
Node(syms.trailer, [Dot(), kw('raises', prefix='')]))
# Adds a 'import pytest' if there wasn't one already
touch_import(None, "pytest", node)
def _modify_power(self, node):
prefix = node.children[0].prefix
# remove old prefix
parts = dotted_parts(self.old_name)
for _ in range((len(parts) + 1) // 2):
node.children.pop(0)
# add new prefix
head = Name(self.new_name.split('.', maxsplit=1)[0], prefix=prefix)
children = []
for part in dotted_parts(self.new_name)[2::2]:
children.append(
Node(
type=syms.trailer,
children=[Dot(), Name(part)],
))
node.children = [head] + children + node.children
def replace_unicode_methods(node, capture, arguments):
# remove any existing __str__ method
b = find_binding("__str__", capture['suite'])
if b and b.type == syms.funcdef:
b.remove()
# rename __unicode__ to __str__
funcname = capture['funcname'].clone()
funcname.value = '__str__'
capture['funcname'].replace(funcname)
# Add a six import
touch_import(None, "six", node)
# Decorate the class with `@six.python_2_unicode_compatible`
classdef = node.clone()
classdef.prefix = ''
decorated = Node(
syms.decorated,
[
Node(
syms.decorator,
[
Leaf(TOKEN.AT, '@', prefix=node.prefix),
Node(
syms.dotted_name,
[
Name('six'),
Dot(),
Name('python_2_unicode_compatible')
],
),
Newline(),
],
prefix=node.prefix,
),
classdef,
],
prefix=node.prefix,
)
node.replace(decorated)
def handle_assertraises(node, capture, arguments):
"""
with self.assertRaises(x):
--> with pytest.raises(x):
self.assertRaises(ValueError, func, arg1)
--> pytest.raises(ValueError, func, arg1)
"""
capture["self_attr"].replace(
keyword("pytest", prefix=capture["self_attr"].prefix))
capture["raises_attr"].replace(
Node(syms.trailer, [Dot(), keyword("raises", prefix="")]))
# Let's remove the msg= keyword argument if found
for child in node.children:
if child.type != syms.trailer:
continue
for tchild in child.children:
if tchild.type != syms.arglist:
continue
previous_argument = None
for argument in tchild.children:
if isinstance(argument, Leaf):
previous_argument = argument
continue
if isinstance(argument, Node):
if argument.type != syms.argument:
previous_argument = argument
continue
for leaf in argument.leaves():
if leaf.value == "msg":
argument.remove()
if previous_argument.value == ",":
# previous_argument is a comma, remove it.
previous_argument.remove()
# Adds a 'import pytest' if there wasn't one already
touch_import(None, "pytest", node)
def encapsulate_transform(
node: LN, capture: Capture, filename: Filename
) -> None:
if "attr_assignment" in capture:
leaf = capture["attr_name"]
leaf.replace(Name(new_name, prefix=leaf.prefix))
if make_property:
# TODO: capture and use type annotation from original assignment
class_node = get_class(node)
suite = find_first(class_node, SYMBOL.suite)
assert isinstance(suite, Node)
indent_node = find_first(suite, TOKEN.INDENT)
assert isinstance(indent_node, Leaf)
indent = indent_node.value
getter = Node(
SYMBOL.decorated,
[
Node(
SYMBOL.decorator,
[
Leaf(TOKEN.AT, "@"),
Name("property"),
Leaf(TOKEN.NEWLINE, "\n"),
],
),
Node(
SYMBOL.funcdef,
[
Name("def", indent),
Name(old_name, prefix=" "),
Node(
SYMBOL.parameters,
[LParen(), Name("self"), RParen()],
),
Leaf(TOKEN.COLON, ":"),
Node(
SYMBOL.suite,
[
Newline(),
Leaf(TOKEN.INDENT, indent.value + " "),
Node(
SYMBOL.simple_stmt,
[
Node(
SYMBOL.return_stmt,
[
Name("return"),
Node(
SYMBOL.power,
Attr(
Name("self"),
Name(new_name),
),
prefix=" ",
),
],
),
Newline(),
],
),
Leaf(TOKEN.DEDENT, "\n" + indent),
],
),
],
prefix=indent,
),
],
)
setter = Node(
SYMBOL.decorated,
[
Node(
SYMBOL.decorator,
[
Leaf(TOKEN.AT, "@"),
Node(
SYMBOL.dotted_name,
[Name(old_name), Dot(), Name("setter")],
),
Leaf(TOKEN.NEWLINE, "\n"),
],
),
Node(
SYMBOL.funcdef,
[
Name("def", indent),
Name(old_name, prefix=" "),
Node(
SYMBOL.parameters,
[
LParen(),
Node(
SYMBOL.typedargslist,
[
Name("self"),
Comma(),
Name("value", prefix=" "),
],
),
RParen(),
],
),
Leaf(TOKEN.COLON, ":"),
Node(
SYMBOL.suite,
[
Newline(),
Leaf(TOKEN.INDENT, indent + " "),
Node(
SYMBOL.simple_stmt,
[
Node(
SYMBOL.expr_stmt,
[
Node(
SYMBOL.power,
Attr(
Name("self"),
Name(new_name),
),
),
Leaf(
TOKEN.EQUAL,
"=",
prefix=" ",
),
Name("value", prefix=" "),
],
),
Newline(),
],
),
Leaf(TOKEN.DEDENT, "\n" + indent),
],
),
],
prefix=indent,
),
],
)
suite.insert_child(-1, getter)
suite.insert_child(-1, setter)
prev = find_previous(getter, TOKEN.DEDENT, recursive=True)
curr = find_last(setter, TOKEN.DEDENT, recursive=True)
assert isinstance(prev, Leaf) and isinstance(curr, Leaf)
prev.prefix, curr.prefix = curr.prefix, prev.prefix
prev.value, curr.value = curr.value, prev.value
def handle_assertraises_regex(node, capture, arguments):
"""
with self.assertRaisesRegex(x, "regex match"):
--> with pytest.raises(x, match="regex match"):
self.assertRaises(ValueError, "regex match", func, arg1)
--> pytest.raises(ValueError, func, arg1, match="regex match")
"""
capture["self_attr"].replace(
keyword("pytest", prefix=capture["self_attr"].prefix))
capture["raises_attr"].replace(
Node(syms.trailer, [Dot(), keyword("raises", prefix="")]))
# Let's remove the msg= keyword argument if found and replace the second argument
# with a match keyword argument
regex_match = None
for child in node.children:
if child.type != syms.trailer:
continue
for tchild in child.children:
if tchild.type != syms.arglist:
continue
previous_argument = None
argnum = 0
for argument in list(tchild.children):
if isinstance(argument, Leaf):
if argument.value != ",":
argnum += 1
else:
previous_argument = argument
continue
if argnum != 2:
previous_argument = argument
continue
if argnum == 2:
regex_match = Node(
syms.argument,
[
Leaf(TOKEN.NAME, "match"),
Leaf(TOKEN.EQUAL, "="),
Leaf(TOKEN.STRING, argument.value),
],
prefix=" ",
)
argument.remove()
if previous_argument and previous_argument.value == ",":
previous_argument.remove()
previous_argument = None
continue
if isinstance(argument, Node):
if argument.type != syms.argument:
previous_argument = argument
continue
for leaf in argument.leaves():
if leaf.value == "msg":
argument.remove()
if previous_argument and previous_argument.value == ",":
# previous_argument is a comma, remove it.
previous_argument.remove()
if regex_match:
regex_match_added = False
for child in node.children:
if regex_match_added:
break
if child.type != syms.trailer:
continue
for tchild in child.children:
if tchild.type != syms.arglist:
continue
tchild.children.append(Leaf(TOKEN.COMMA, ","))
tchild.children.append(regex_match)
regex_match_added = True
break
# Adds a 'import pytest' if there wasn't one already
touch_import(None, "pytest", node)
def make_pytest_raises_blocks(node, capture, filename):
"""
Turns this:
try:
...
pytest.fail(...)
except:
pass
Into:
with pytest.raises(Exception):
...
Not only is this prettier, but the former is a bug since
pytest.fail() raises an exception.
"""
exc_class = capture.get('exc_class', None)
if exc_class:
exc_class = exc_class.clone()
exc_class.prefix = ''
raises_args = [exc_class]
else:
raises_args = [kw('Exception', prefix='')]
reason = capture.get('reason')
if reason:
assert len(reason) == 1
reason = KeywordArg(kw('message'), reason[0].clone())
raises_args = [Node(syms.arglist, raises_args + [Comma(), reason])]
raises_args = [LParen()] + raises_args + [RParen()]
capture['fail_stmt'].remove()
try_suite = capture['try_suite'].clone()
with_stmt = Node(
syms.with_stmt,
[
kw('with', prefix=''),
Node(
syms.power,
[
kw('pytest'),
Node(syms.trailer, [Dot(), kw('raises', prefix='')]),
Node(syms.trailer, raises_args),
],
),
Leaf(TOKEN.COLON, ':'),
try_suite,
],
prefix=node.prefix,
)
# Trailing whitespace and any comments after the if statement are captured
# in the prefix for the dedent node. Copy it to the following node.
dedent = capture["dedent"]
next_node = node.next_sibling
# This extra newline avoids syntax errors in some cases (where the try
# statement is at the end of another suite)
# I don't really know why those occur.
# Should clean this stuff up with `black` later.
node.replace([with_stmt, Newline()])
next_node.prefix = dedent.prefix