def _make_bare_import_node(name: str, trailing_nl: bool = False) -> Node:
assert name # non-empty
children = [
Leaf(token.NAME, "import"),
Leaf(token.NAME, name, prefix=" "),
Newline(),
]
if trailing_nl:
children.append(Newline())
return Node(
syms.import_name,
children,
)
def _make_from_import_node(left: str,
right: Sequence[str],
trailing_nl: bool = False) -> Node:
assert right # non-empty
name_leaves = [Leaf(token.NAME, right[0], prefix=" ")]
name_leaves.extend(
Leaf(token.NAME, name, prefix=", ") for name in right[1:])
children = [
Leaf(token.NAME, "from"),
Leaf(token.NAME, left, prefix=" "),
Leaf(token.NAME, "import", prefix=" "),
Node(syms.import_as_names, name_leaves),
Newline(),
]
if trailing_nl:
children.append(Newline())
return Node(syms.import_from, children)
def newline_node(node):
"""
Find a NEWLINE node from AST tree, and clone it to create a new NEWLINE node.
So we don't need to care about os platform.
if no NEWLINE node found in AST tree, create a new NEWLINE node according to os platform.
"""
global _newline_node
if _newline_node is None:
p = node
while p is not None and p.parent is not None:
p = p.parent
if p is not None:
for leaf in p.leaves():
if leaf.type == token.NEWLINE:
_newline_node = leaf
break
if _newline_node is None:
if is_windows:
_newline_node = Newline(value="\r\n")
else:
_newline_node = Newline()
return _newline_node.clone()
def remove_success_expectations(node, capture, filename):
"""
We're about to remove the 'success' return value of all the helpers,
and just let them pytest.fail() themselves.
So we need to remove where tests are expecting them to return 'success'.
if x() != 'success':
return 'fail'
--> x()
"""
if capture['returntype'].type not in (TOKEN.STRING, TOKEN.NAME):
return
if capture['x'].type == TOKEN.NAME:
# `if ret == 'success'`.
# We can just remove this.
if (capture['returntype'] == TOKEN.NAME
and capture['returntype'].value != capture['x'].value):
# not sure what this is doing? leave it alone.
return
# 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
node.remove()
next_node.prefix = dedent.prefix
elif capture['x'].type == syms.power:
# is a function call. call it, just discard the result.
if capture['returntype'].type == TOKEN.NAME:
# not sure what this is doing? leave it alone
return
func = capture['x'].clone()
func.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
node.replace([func, Newline()])
next_node.prefix = dedent.prefix
else:
print(capture['x'])
raise ValueError("unknown type")
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 replace_if_with_assert(if_node, dedent_node, assert_node):
# Trailing whitespace and any comments after the if statement are captured
# in the prefix for the dedent node. Copy it to the following node.
next_node = if_node.next_sibling
if_node.replace([assert_node, Newline()])
next_node.prefix = dedent_node.prefix
message_nodes = assert_node.children[3:]
if message_nodes:
# If one line, and it's longer than LONG_LINE, make a basic attempt to split before the message
lines = str(assert_node).splitlines()
original_num_long_lines = sum(1 for line in lines
if len(line) > LONG_LINE)
if original_num_long_lines:
# First, clone the original.
assert_node, split_assert = assert_node.clone(), assert_node
# now insert a newline after the comma
reason_node = message_nodes[0]
reason_node.prefix = f" \\\n{find_indentation(split_assert)} {reason_node.prefix.lstrip(' ')}"
def transform(self, node, results):
# First, find the sys import. We'll just hope it's global scope.
if "sys_import" in results:
if self.sys_import is None:
self.sys_import = results["sys_import"]
return
func = results["func"].clone()
func.prefix = ""
register = pytree.Node(syms.power, Attr(Name("atexit"), Name("register")))
call = Call(register, [func], node.prefix)
node.replace(call)
if self.sys_import is None:
# That's interesting.
self.warning(
node,
"Can't find sys import; Please add an atexit "
"import at the top of your file.",
)
return
# Now add an atexit import after the sys import.
names = self.sys_import.children[1]
if names.type == syms.dotted_as_names:
names.append_child(Comma())
names.append_child(Name("atexit", " "))
else:
containing_stmt = self.sys_import.parent
position = containing_stmt.children.index(self.sys_import)
stmt_container = containing_stmt.parent
new_import = pytree.Node(
syms.import_name, [Name("import"), Name("atexit", " ")]
)
new = pytree.Node(syms.simple_stmt, [new_import])
containing_stmt.insert_child(position + 1, Newline())
containing_stmt.insert_child(position + 2, new)
def transform_member(self, node, results):
"""Transform for imports of specific module elements. Replaces
the module to be imported from with the appropriate new
module.
"""
mod_member = results.get("mod_member")
pref = mod_member.prefix
member = results.get("member")
# Simple case with only a single member being imported
if member:
# this may be a list of length one, or just a node
if isinstance(member, list):
member = member[0]
new_name = None
for change in MAPPING[mod_member.value]:
if member.value in change[1]:
new_name = change[0]
break
if new_name:
mod_member.replace(Name(new_name, prefix=pref))
else:
self.cannot_convert(node, "This is an invalid module element")
# Multiple members being imported
else:
# a dictionary for replacements, order matters
modules = []
mod_dict = {}
members = results["members"]
for member in members:
# we only care about the actual members
if member.type == syms.import_as_name:
as_name = member.children[2].value
member_name = member.children[0].value
else:
member_name = member.value
as_name = None
if member_name != ",":
for change in MAPPING[mod_member.value]:
if member_name in change[1]:
if change[0] not in mod_dict:
modules.append(change[0])
mod_dict.setdefault(change[0], []).append(member)
new_nodes = []
indentation = find_indentation(node)
first = True
def handle_name(name, prefix):
if name.type == syms.import_as_name:
kids = [
Name(name.children[0].value, prefix=prefix),
name.children[1].clone(),
name.children[2].clone(),
]
return [Node(syms.import_as_name, kids)]
return [Name(name.value, prefix=prefix)]
for module in modules:
elts = mod_dict[module]
names = []
for elt in elts[:-1]:
names.extend(handle_name(elt, pref))
names.append(Comma())
names.extend(handle_name(elts[-1], pref))
new = FromImport(module, names)
if not first or node.parent.prefix.endswith(indentation):
new.prefix = indentation
new_nodes.append(new)
first = False
if new_nodes:
nodes = []
for new_node in new_nodes[:-1]:
nodes.extend([new_node, Newline()])
nodes.append(new_nodes[-1])
node.replace(nodes)
else:
self.cannot_convert(node, "All module elements are invalid")
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 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