def transform(self, node, results):
suite = results['suite'].clone()
# todo: handle tabs
dedent = len(find_indentation(suite)) - len(find_indentation(node))
self.dedent(suite, dedent)
# remove the first newline behind the classdef header
first = suite.children[0]
if first.type == token.NEWLINE:
if len(first.value) == 1:
del suite.children[0]
else:
first.value == first.value[1:]
return suite
def transform(self, node, results):
suite = results['suite'].clone()
# todo: handle tabs
dedent = len(find_indentation(suite)) - len(find_indentation(node))
self.dedent(suite, dedent)
# remove the first newline behind the classdef header
first = suite.children[0]
if first.type == token.NEWLINE:
if len(first.value) == 1:
del suite.children[0]
else:
first.value == first.value[1:]
return suite
def add_py2_annot(self, argtypes, restype, node, results):
children = results['suite'][0].children
# Insert '# type: {annot}' comment.
# For reference, see lib2to3/fixes/fix_tuple_params.py in stdlib.
if len(children) >= 1 and children[0].type != token.NEWLINE:
# one liner function
if children[0].prefix.strip() == '':
children[0].prefix = ''
children.insert(0, Leaf(token.NEWLINE, '\n'))
children.insert(
1, Leaf(token.INDENT,
find_indentation(node) + ' '))
children.append(Leaf(token.DEDENT, ''))
if len(children) >= 2 and children[1].type == token.INDENT:
degen_str = '(...) -> %s' % restype
short_str = '(%s) -> %s' % (', '.join(argtypes), restype)
if (len(short_str) > 64
or len(argtypes) > 5) and len(short_str) > len(degen_str):
self.insert_long_form(node, results, argtypes)
annot_str = degen_str
else:
annot_str = short_str
children[1].prefix = '%s# type: %s\n%s' % (
children[1].value, annot_str, children[1].prefix)
children[1].changed()
else:
self.log_message(
"%s:%d: cannot insert annotation for one-line function" %
(self.filename, node.get_lineno()))
def transform(self, node, results):
# Sometimes newlines are in prefix of current node, sometimes they're
# in prefix of the prev sibling
if node.prefix.count('\n'):
newline_node = node
else:
newline_node = get_whitespace_before_definition(node)
if not newline_node:
# No previous node, must be the first node.
return
if newline_node.type in [token.INDENT, token.NEWLINE]:
# If the newline_node is an indent or newline, we don't need to
# worry about fixing indentation since it is not part of the
# prefix. Dedents do have it as part of the prefix.
curr_node_indentation = ''
else:
curr_node_indentation = find_indentation(node)
min_lines_between_defs, max_lines_between_defs = (self.
get_newline_limits(node))
new_prefix = self.trim_comments(curr_node_indentation, newline_node.
prefix, min_lines_between_defs, max_lines_between_defs)
if newline_node.prefix != new_prefix:
newline_node.prefix = new_prefix
newline_node.changed()
def transform(self, node, results):
def process_arg(arg):
if isinstance(arg, Leaf) and arg.type == token.COMMA:
return
elif isinstance(arg, Node) and arg.type == syms.argument:
# keyword argument
name, equal, value = arg.children
assert name.type == token.NAME # what is the symbol for 1?
assert equal.type == token.EQUAL # what is the symbol for 1?
value = value.clone()
value.prefix = " "
kwargs[name.value] = value
else:
assert not kwargs, 'all positional args are assumed to come first'
posargs.append(arg.clone())
method = results['method'][0].value
# map (deprecated) aliases to original to avoid analysing
# the decorator function
method = _method_aliases.get(method, method)
posargs = []
kwargs = {}
# This is either a "arglist" or a single argument
if results['arglist'].type == syms.arglist:
for arg in results['arglist'].children:
process_arg(arg)
else:
process_arg(results['arglist'])
try:
test_func = getattr(unittest.TestCase, method)
except AttributeError:
raise RuntimeError("Your unittest package does not support '%s'. "
"consider updating the package" % method)
required_args, argsdict = utils.resolve_func_args(test_func, posargs, kwargs)
if method.startswith(('assertRaises', 'assertWarns')):
n_stmt = _method_map[method](*required_args,
indent=find_indentation(node),
kws=argsdict,
arglist=results['arglist'])
else:
n_stmt = Node(syms.assert_stmt,
[Name('assert'),
_method_map[method](*required_args, kws=argsdict)])
if argsdict.get('msg', None) is not None:
n_stmt.children.extend((Name(','), argsdict['msg']))
n_stmt.prefix = node.prefix
return n_stmt
def transform(self, node, results):
def process_arg(arg):
if isinstance(arg, Leaf) and arg.type == token.COMMA:
return
elif isinstance(arg, Node) and arg.type == syms.argument:
# keyword argument
name, equal, value = arg.children
assert name.type == token.NAME # what is the symbol for 1?
assert equal.type == token.EQUAL # what is the symbol for 1?
value = value.clone()
value.prefix = " "
kwargs[name.value] = value
else:
assert not kwargs, 'all positional args are assumed to come first'
posargs.append(arg.clone())
method = results['method'][0].value
# map (deprecated) aliases to original to avoid analysing
# the decorator function
method = _method_aliases.get(method, method)
posargs = []
kwargs = {}
# This is either a "arglist" or a single argument
if results['arglist'].type == syms.arglist:
for arg in results['arglist'].children:
process_arg(arg)
else:
process_arg(results['arglist'])
try:
test_func = getattr(unittest.TestCase, method)
except AttributeError:
raise RuntimeError("Your unittest package does not support '%s'. "
"consider updating the package" % method)
required_args, argsdict = utils.resolve_func_args(test_func, posargs, kwargs)
if method.startswith(('assertRaises', 'assertWarns')):
n_stmt = _method_map[method](*required_args,
indent=find_indentation(node),
kws=argsdict,
arglist=results['arglist'])
else:
n_stmt = Node(syms.assert_stmt,
[Name('assert'),
_method_map[method](*required_args, kws=argsdict)])
if argsdict.get('msg', None) is not None:
n_stmt.children.extend((Name(','), argsdict['msg']))
n_stmt.prefix = node.prefix
return n_stmt
def find_indentation(node):
try:
from lib2to3.fixer_util import find_indentation
return find_indentation(node)
except ImportError:
while node is not None:
if node.type == symbols.suite and len(node.children) > 2:
indent = node.children[1]
if indent.type == token.INDENT:
return indent.value
node = node.parent
return ""
def find_indentation(node):
try:
from lib2to3.fixer_util import find_indentation
return find_indentation(node)
except ImportError:
while node is not None:
if node.type == symbols.suite and len(node.children) > 2:
indent = node.children[1]
if indent.type == token.INDENT:
return indent.value
node = node.parent
return ""
def transform(self, node, results):
unifunc = results["unifunc"]
strfunc = Name("__str__", prefix=unifunc.prefix)
unifunc.replace(strfunc)
klass = node.clone()
klass.prefix = '\n' + find_indentation(node)
decorator = Node(syms.decorator, [Leaf(token.AT, "@"), Name('python_2_unicode_compatible')])
decorated = Node(syms.decorated, [decorator, klass], prefix=node.prefix)
node.replace(decorated)
touch_import('django.utils.encoding', 'python_2_unicode_compatible', decorated)
def transform_semi(self, node):
for child in node.children:
if child.type == token.SEMI:
# Strip any whitespace from the next sibling
if child.next_sibling.prefix != child.next_sibling.prefix.lstrip():
child.next_sibling.prefix = child.next_sibling.prefix.lstrip()
child.next_sibling.changed()
# Replace the semi with a newline
old_depth = find_indentation(child)
child.replace([Leaf(token.NEWLINE, u"\n"), Leaf(token.INDENT, old_depth)])
child.changed()
return node
def transform_colon(self, node):
node_copy = node.clone()
# Strip any whitespace that could have been there
node_copy.prefix = node_copy.prefix.lstrip()
old_depth = find_indentation(node)
new_indent = u'%s%s' % ((u' ' * 4), old_depth)
new_node = Node(symbols.suite, [Leaf(token.NEWLINE, u'\n'), Leaf(token
.INDENT, new_indent), node_copy, Leaf(token.DEDENT, u'')])
node.replace(new_node)
node.changed()
# Replace node with new_node in case semi
return node_copy
def transform(self, node, results):
if "fix_print_with_import" in node.prefix:
return node
r = super(FixPrintWithImport, self).transform(node, results)
if not r or r == node:
return r
if not r.prefix:
indentation = find_indentation(node)
r.prefix = "# fix_print_with_import\n" + indentation
else:
r.prefix = re.sub('([ \t]*$)', r'\1# fix_print_with_import\n\1', r.prefix)
return r
def transform(self, node, results):
unifunc = results["unifunc"]
strfunc = Name("__str__", prefix=unifunc.prefix)
unifunc.replace(strfunc)
klass = node.clone()
klass.prefix = '\n' + find_indentation(node)
decorator = Node(
syms.decorator,
[Leaf(token.AT, "@"),
Name('python_2_unicode_compatible')])
decorated = Node(syms.decorated, [decorator, klass],
prefix=node.prefix)
node.replace(decorated)
touch_import('django.utils.encoding', 'python_2_unicode_compatible',
decorated)
def transform(self, node, results):
if "fix_print_with_import" in node.prefix:
return node
r = super(FixPrintWithImport, self).transform(node, results)
if not r or r == node:
return r
if not r.prefix:
indentation = find_indentation(node)
r.prefix = "# fix_print_with_import\n" + indentation
else:
r.prefix = re.sub('([ \t]*$)', r'\1# fix_print_with_import\n\1',
r.prefix)
return r
def patch(t):
visitor = Visitor()
visitor.visit(t)
seen = set()
for node in visitor.r:
target = node
while target:
if node_name(target) == "funcdef":
if id(target) not in seen:
seen.add(id(target)) # xxx:
insert_before(target, Decorator("profile"))
target.prefix = u.find_indentation(target)
elif target.parent is None:
break
target = target.parent
def _insert_node(t0, node) -> None:
ancestors = []
target = node
while target is not None:
ancestors.append(node_name(target))
target = target.parent
assert ancestors[-1] == "file_input"
is_toplevel_def = ancestors[-2] in ("funcdef", "classdef", "async_funcdef")
if not is_toplevel_def:
return
indentation = find_indentation(node)
if indentation == "":
node.prefix = "\n\n"
else:
node.prefix = f"\n{indentation}"
t0.append_child(node)
def transform(self, node, results):
child_imports = [leaf.value for leaf in results if leaf.type == token.
NAME]
current_indentation = find_indentation(node)
new_nodes = []
for index, module_name in enumerate(child_imports):
new_prefix = current_indentation
if not index:
# Don't add more indentation if this is the first one
new_prefix = None
new_nodes.append(Node(symbols.simple_stmt, [Node(symbols.
import_name, [Leaf(token.NAME, 'import', prefix=new_prefix),
Leaf(token.NAME, module_name, prefix=" ")]), Leaf(token.
NEWLINE, '\n')]))
node.replace(new_nodes)
node.changed()
def transform(self, node: LN, capture: Capture) -> None:
imports = []
for n in capture["module_names"]:
if n.type == token.COMMA:
pass
elif n.type == token.NAME:
imports.append((n.value, None))
elif n.type == syms.dotted_name:
imports.append((traverse_dotted_name(n), None))
elif n.type == syms.dotted_as_name:
import_name, import_nick = n.children[0], n.children[2]
assert n.children[1].type == token.NAME and n.children[
1].value == "as"
imports.append(
(traverse_dotted_name(import_name), import_nick.value))
rename_seen = False
nodes = []
indent = find_indentation(node)
for name, nick in imports:
try:
new_name = get_new_name(name)
except NameRemovedError as exc:
self.warn(node, str(exc))
continue
if new_name:
rename_seen = True
name = new_name
new_node = Node(
syms.import_name,
[Name("import"),
ImportAsName(name, nick, prefix=" ")],
prefix=f"\n{indent}",
)
nodes.append(new_node)
if not nodes:
return
nodes[0].prefix = node.prefix
if rename_seen:
node.replace(nodes)
def add_py2_annot(self, argtypes, restype, node, results):
# type: (List[str], str, Node, Dict[str, Any]) -> None
children = results['suite'][0].children
# Insert '# type: {annot}' comment.
# For reference, see lib2to3/fixes/fix_tuple_params.py in stdlib.
if len(children) >= 1 and children[0].type != token.NEWLINE:
# one liner function
if children[0].prefix.strip() == '':
children[0].prefix = ''
children.insert(0, Leaf(token.NEWLINE, '\n'))
children.insert(
1, Leaf(token.INDENT, find_indentation(node) + ' '))
children.append(Leaf(token.DEDENT, ''))
if len(children) >= 2 and children[1].type == token.INDENT:
degen_str = '(...) -> %s' % restype
short_str = '(%s) -> %s' % (', '.join(argtypes), restype)
if self.use_py2_long_form(argtypes, short_str, degen_str):
self.insert_long_form(node, results, argtypes)
annot_str = degen_str
else:
annot_str = short_str
indent_node = children[1]
comment, sep, other_comments = indent_node.prefix.partition('\n')
comment = comment.rstrip() + sep
annot_str = '# type: %s\n' % (annot_str,)
if comment == annot_str:
return
if comment and not is_type_comment(comment):
# push existing non-type comment to next line
annot_str += comment
indent_node.prefix = indent_node.value + annot_str + other_comments
indent_node.changed()
else:
self.log_message("%s:%d: cannot insert annotation for one-line function" %
(self.filename, node.get_lineno()))
def transform(self, node, results):
# Determine the node's column number by finding the first leaf.
leaf = node
while not isinstance(leaf, Leaf):
leaf = leaf.children[0]
# Only match functions and the global indentation level.
if leaf.column != 0:
return
indent = None
for child in node.children:
if isinstance(child, Node) and child.type == python_symbols.suite:
indent = find_indentation(child)
if isinstance(
child, Leaf
) and child.type == token.NAME and child.value == self.funcname:
child.value = self.newname
elif isinstance(child,
Node) and child.type == python_symbols.parameters:
pre_params = []
for param in self.pre_params:
pre_params.append(Leaf(token.NAME, param))
pre_params.append(Leaf(token.COMMA, ', '))
child.children[1:1] = pre_params
post_params = []
for param in self.post_params:
post_params.append(Leaf(token.COMMA, ','))
post_params.append(Leaf(token.NAME, param))
child.children[-1:-1] = post_params
if child.children[-2].type == token.COMMA:
child.children.pop(-2)
child.changed()
if self.add_statement:
node.children.append(
Leaf(0, indent + self.add_statement.rstrip() + '\n'))
if self.remove:
self.results.append(node)
node.replace([])
return None
else:
return node
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 != u",":
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 transform(self, node, results):
if FixAnnotate.counter is not None:
if FixAnnotate.counter <= 0:
return
# Check if there's already a long-form annotation for some argument.
parameters = results.get('parameters')
if parameters is not None:
for ch in parameters.pre_order():
if ch.prefix.lstrip().startswith('# type:'):
return
args = results.get('args')
if args is not None:
for ch in args.pre_order():
if ch.prefix.lstrip().startswith('# type:'):
return
suite = results['suite']
children = suite[0].children
# NOTE: I've reverse-engineered the structure of the parse tree.
# It's always a list of nodes, the first of which contains the
# entire suite. Its children seem to be:
#
# [0] NEWLINE
# [1] INDENT
# [2...n-2] statements (the first may be a docstring)
# [n-1] DEDENT
#
# Comments before the suite are part of the INDENT's prefix.
#
# "Compact" functions (e.g. "def foo(x, y): return max(x, y)")
# have a different structure (no NEWLINE, INDENT, or DEDENT).
# Check if there's already an annotation.
for ch in children:
if ch.prefix.lstrip().startswith('# type:'):
return # There's already a # type: comment here; don't change anything.
# Compute the annotation
annot = self.make_annotation(node, results)
if annot is None:
return
# Insert '# type: {annot}' comment.
# For reference, see lib2to3/fixes/fix_tuple_params.py in stdlib.
if len(children) >= 1 and children[0].type != token.NEWLINE:
if children[0].prefix.strip() == '':
children[0].prefix = ''
children.insert(0, Leaf(token.NEWLINE, '\n'))
children.insert(
1, Leaf(token.INDENT,
find_indentation(node) + ' '))
children.append(Leaf(token.DEDENT, ''))
if len(children) >= 2 and children[1].type == token.INDENT:
argtypes, restype = annot
degen_str = '(...) -> %s' % restype
short_str = '(%s) -> %s' % (', '.join(argtypes), restype)
if (len(short_str) > 64
or len(argtypes) > 5) and len(short_str) > len(degen_str):
self.insert_long_form(node, results, argtypes)
annot_str = degen_str
else:
annot_str = short_str
children[1].prefix = '%s# type: %s\n%s' % (
children[1].value, annot_str, children[1].prefix)
children[1].changed()
if FixAnnotate.counter is not None:
FixAnnotate.counter -= 1
# Also add 'from typing import Any' at the top if needed.
self.patch_imports(argtypes + [restype], node)
else:
self.log_message(
"%s:%d: cannot insert annotation for one-line function" %
(self.filename, node.get_lineno()))
def transform(self, node, results):
def process_arg(arg):
if isinstance(arg, Leaf) and arg.type == token.COMMA:
return
elif isinstance(arg, Node) and arg.type == syms.argument:
# keyword argument
name, equal, value = arg.children
assert name.type == token.NAME # what is the symbol for 1?
assert equal.type == token.EQUAL # what is the symbol for 1?
value = value.clone()
value.prefix = " "
kwargs[name.value] = value
else:
assert not kwargs, 'all positional args are assumed to come first'
posargs.append(arg.clone())
func = results['func'].value
custom_helper = False
if node.parent.type == syms.return_stmt:
# custom helper with `return eq_(...)`
# We're not rendering the `assert` in that case
# to allow the code to continue functioning
custom_helper = True
posargs = []
kwargs = []
if results['arglist'].type == syms.arglist:
for arg in results['arglist'].children:
process_arg(arg)
else:
process_arg(results['arglist'])
if len(posargs) == 2:
left, right = posargs
elif len(posargs) == 3:
left, right, _ = posargs
left.prefix = " "
right.prefix = " "
strip_newlines(left)
strip_newlines(right)
# Ignore customized assert messages for now
if isinstance(right,
Leaf) and right.value in ('None', 'True', 'False'):
op = Name('is', prefix=' ')
body = [Node(syms.comparison, (left, op, right))]
else:
op = Name('==', prefix=' ')
body = [Node(syms.comparison, (left, op, right))]
indent = find_indentation(node)
ret = Name('assert')
if node.parent.prefix.endswith(indent):
ret.prefix = indent
if custom_helper:
return body
return [ret] + body
def is_top_level(node):
"""Is node at top indentation level (i.e. module globals)?"""
return bool(len(find_indentation(node)))
def is_top_level(node):
"""Is node at top indentation level (module globals)"""
return 0 == len(find_indentation(node))
def transform(self, node, results):
def process_arg(arg):
if isinstance(arg, Leaf) and arg.type == token.COMMA:
return
elif (isinstance(arg, Node) and arg.type == syms.argument
and arg.children[1].type == token.EQUAL):
# keyword argument
name, equal, value = arg.children
assert name.type == token.NAME
assert equal.type == token.EQUAL
value = value.clone()
kwargs[name.value] = value
if '\n' in arg.prefix:
value.prefix = arg.prefix
else:
value.prefix = arg.prefix.strip() + " "
else:
if (isinstance(arg, Node) and arg.type == syms.argument
and arg.children[0].type == 36
and arg.children[0].value == '**'):
return
assert not kwargs, 'all positional args are assumed to come first'
if (isinstance(arg, Node) and arg.type == syms.argument
and arg.children[1].type == syms.comp_for):
# argument is a generator expression w/o
# parenthesis, add parenthesis
value = arg.clone()
value.children.insert(0, Leaf(token.LPAR, '('))
value.children.append(Leaf(token.RPAR, ')'))
posargs.append(value)
else:
posargs.append(arg.clone())
method = results['method'][0].value
# map (deprecated) aliases to original to avoid analysing
# the decorator function
method = _method_aliases.get(method, method)
posargs = []
kwargs = {}
# This is either a "arglist" or a single argument
if results['arglist'].type == syms.arglist:
for arg in results['arglist'].children:
process_arg(arg)
else:
process_arg(results['arglist'])
try:
test_func = getattr(unittest.TestCase, method)
except AttributeError:
raise RuntimeError("Your unittest package does not support '%s'. "
"consider updating the package" % method)
required_args, argsdict = utils.resolve_func_args(
test_func, posargs, kwargs)
if method.startswith(('assertRaises', 'assertWarns')):
n_stmt = _method_map[method](*required_args,
indent=find_indentation(node),
kws=argsdict,
arglist=results['arglist'],
node=node)
else:
n_stmt = Node(syms.assert_stmt, [
Name('assert'), _method_map[method](*required_args,
kws=argsdict)
])
if argsdict.get('msg', None) is not None:
n_stmt.children.extend((Name(','), argsdict['msg']))
def fix_line_wrapping(x):
for c in x.children:
# no need to worry about wrapping of "[", "{" and "("
if c.type in [token.LSQB, token.LBRACE, token.LPAR]:
break
if c.prefix.startswith('\n'):
c.prefix = c.prefix.replace('\n', ' \\\n')
fix_line_wrapping(c)
fix_line_wrapping(n_stmt)
# the prefix should be set only after fixing line wrapping because it can contain a '\n'
n_stmt.prefix = node.prefix
# add necessary imports
if 'Raises' in method or 'Warns' in method:
add_import('pytest', node)
if ('Regex' in method and not 'Raises' in method
and not 'Warns' in method):
add_import('re', node)
return n_stmt
def transform(self, node, results):
def process_arg(arg):
if isinstance(arg, Leaf) and arg.type == token.COMMA:
return
elif isinstance(arg, Node) and arg.type == syms.argument:
# keyword argument
name, equal, value = arg.children
assert name.type == token.NAME # what is the symbol for 1?
assert equal.type == token.EQUAL # what is the symbol for 1?
value = value.clone()
kwargs[name.value] = value
if '\n' in arg.prefix:
value.prefix = arg.prefix
else:
value.prefix = arg.prefix.strip() + " "
else:
assert not kwargs, 'all positional args are assumed to come first'
posargs.append(arg.clone())
method = results['method'][0].value
# map (deprecated) aliases to original to avoid analysing
# the decorator function
method = _method_aliases.get(method, method)
posargs = []
kwargs = {}
# This is either a "arglist" or a single argument
if results['arglist'].type == syms.arglist:
for arg in results['arglist'].children:
process_arg(arg)
else:
process_arg(results['arglist'])
try:
test_func = getattr(unittest.TestCase, method)
except AttributeError:
raise RuntimeError("Your unittest package does not support '%s'. "
"consider updating the package" % method)
required_args, argsdict = utils.resolve_func_args(test_func, posargs, kwargs)
if method.startswith(('assertRaises', 'assertWarns')):
n_stmt = _method_map[method](*required_args,
indent=find_indentation(node),
kws=argsdict,
arglist=results['arglist'],
node=node)
else:
n_stmt = Node(syms.assert_stmt,
[Name('assert'),
_method_map[method](*required_args, kws=argsdict)])
if argsdict.get('msg', None) is not None:
n_stmt.children.extend((Name(','), argsdict['msg']))
def fix_line_wrapping(x):
for c in x.children:
# no need to worry about wrapping of "[", "{" and "("
if c.type in [token.LSQB, token.LBRACE, token.LPAR]:
break
if c.prefix.startswith('\n'):
c.prefix = c.prefix.replace('\n', ' \\\n')
fix_line_wrapping(c)
fix_line_wrapping(n_stmt)
# the prefix should be set only after fixing line wrapping because it can contain a '\n'
n_stmt.prefix = node.prefix
# add necessary imports
if 'Raises' in method or 'Warns' in method:
add_import('pytest', node)
if 'Regex' in method:
add_import('re', node)
return n_stmt
def transform(self, node: LN, capture: Capture) -> None:
module_name = traverse_dotted_name(capture["module_name"])
to_process = capture["module_imports"]
imports: List[Tuple[str, Optional[str]]] = []
for n in to_process:
if n.type in (token.COMMA, token.LPAR, token.RPAR):
continue
elif n.type == token.STAR:
self.warn(n, "Cannot guarantee * imports are correct.")
imports.append(("*", None))
elif n.type == token.NAME:
imports.append((n.value, None))
elif n.type == syms.import_as_name:
import_name, import_nick = n.children[0], n.children[2]
assert n.children[1].type == token.NAME and n.children[
1].value == "as"
imports.append((import_name.value,
import_nick.value if import_nick else None))
elif n.type == syms.import_as_names:
to_process.extend(n.children)
else:
raise Exception(f"unexpected node {repr(n)}")
imports_by_package: DefaultDict[str, List[Tuple[
str, Optional[str]]]] = collections.defaultdict(list)
for name, nick in imports:
full_name = f"{module_name}.{name}"
try:
new_full_name = get_new_name(full_name) or full_name
except NameRemovedError as exc:
self.warn(node, str(exc))
continue
package, new_name = split_package_and_name(new_full_name)
if name != new_name and nick is None:
nick = name
imports_by_package[package].append((new_name, nick))
indent = find_indentation(node)
nodes = []
for package, imports in sorted(imports_by_package.items(),
key=lambda i: i[0]):
if package:
result = FromImport(package, imports, prefix=f"\n{indent}")
nodes.append(result)
else:
for name, nick in imports:
nodes.append(
Node(
syms.import_name,
[
Name("import"),
ImportAsName(name, nick, prefix=" ")
],
prefix=f"\n{indent}",
))
if not nodes:
return
nodes[0].prefix = node.prefix
node.replace(nodes)
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 f(x):
def g(y):
return y + y
return g(x + 1)
"""
def Decorator(name):
return u.Node(
278, [u.Leaf(token.AT, "@", prefix=None),
u.Name(name),
u.Newline()])
def insert_before(node, new_node):
for i, x in enumerate(node.parent.children):
if x == node:
node.parent.insert_child(i, new_node)
return True
return False
t = parse_string(code)
for defs in run(t):
insert_before(defs[0], Decorator("profile"))
if not defs[0].prefix:
defs[0].prefix = u.find_indentation(defs[0])
break
print(t)
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')
if member:
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')
else:
modules = []
mod_dict = {}
members = results['members']
for member in 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 != u',':
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')
return
def transform(self, node, results):
def process_arg(arg):
if isinstance(arg, Leaf) and arg.type == token.COMMA:
return
elif isinstance(arg, Node) and arg.type == syms.argument:
# keyword argument
name, equal, value = arg.children
assert name.type == token.NAME # what is the symbol for 1?
assert equal.type == token.EQUAL # what is the symbol for 1?
value = value.clone()
value.prefix = " "
kwargs[name.value] = value
else:
assert not kwargs, 'all positional args are assumed to come first'
posargs.append(arg.clone())
func = results['func'].value
custom_helper = False
if node.parent.type == syms.return_stmt:
# custom helper with `return eq_(...)`
# We're not rendering the `assert` in that case
# to allow the code to continue functioning
custom_helper = True
posargs = []
kwargs = []
if results['arglist'].type == syms.arglist:
for arg in results['arglist'].children:
process_arg(arg)
else:
process_arg(results['arglist'])
if len(posargs) == 2:
left, right = posargs
elif len(posargs) == 3:
left, right, _ = posargs
left.prefix = " "
right.prefix = " "
strip_newlines(left)
strip_newlines(right)
# Ignore customized assert messages for now
if isinstance(right, Leaf) and right.value in ('None', 'True', 'False'):
op = Name('is', prefix=' ')
body = [Node(syms.comparison, (left, op, right))]
else:
op = Name('==', prefix=' ')
body = [Node(syms.comparison, (left, op, right))]
indent = find_indentation(node)
ret = Name('assert')
if node.parent.prefix.endswith(indent):
ret.prefix = indent
if custom_helper:
return body
return [ret] + body
def is_top_level(node):
"""Is node at top indentation level (module globals)"""
return 0 == len(find_indentation(node))
def fix_leaves(self, node_to_split):
parent_depth = find_indentation(node_to_split)
new_indent = "%s%s" % (' ' * 4, parent_depth)
# For now, just indent additional lines by 4 more spaces
child_leaves = []
combined_prefix = ""
prev_leaf = None
for index, leaf in enumerate(node_to_split.leaves()):
if index and leaf.prefix.count('#'):
if not combined_prefix:
combined_prefix = "%s#" % new_indent
combined_prefix += leaf.prefix.split('#')[-1]
# We want to strip all newlines so we can properly insert newlines
# where they should be
if leaf.type != token.NEWLINE:
if leaf.prefix.count('\n') and index:
# If the line contains a newline, we need to strip all
# whitespace since there were leading indent spaces
if (prev_leaf and prev_leaf.type in [token.DOT, token.LPAR]
or leaf.type in [token.RPAR]):
leaf.prefix = ""
else:
leaf.prefix = " "
# Append any trailing inline comments to the combined
# prefix
child_leaves.append(leaf)
prev_leaf = leaf
# Like TextWrapper, but for nodes. We split on MAX_CHARS - 1 since we
# may need to insert a leading parenth. It's not great, but it would be
# hard to do properly.
split_leaves = wrap_leaves(child_leaves, width=MAX_CHARS - 1,
subsequent_indent=new_indent)
new_node = Node(node_to_split.type, [])
# We want to keep track of if we are breaking inside a parenth
open_count = 0
need_parens = False
for line_index, curr_line_nodes in enumerate(split_leaves):
for node_index, curr_line_node in enumerate(curr_line_nodes):
if line_index and not node_index:
# If first node in non-first line, reset prefix since there
# may have been spaces previously
curr_line_node.prefix = new_indent
new_node.append_child(curr_line_node)
if curr_line_node.type in OPENING_TOKENS:
open_count += 1
if curr_line_node.type in CLOSING_TOKENS:
open_count -= 1
if line_index != len(split_leaves) - 1:
# Don't add newline at the end since it it part of the next
# sibling
new_node.append_child(Leaf(token.NEWLINE, '\n'))
# Checks if we ended a line without being surrounded by parens
if open_count <= 0:
need_parens = True
if need_parens:
# Parenthesize the parent if we're not inside parenths, braces,
# brackets, since we inserted newlines between leaves
parenth_before_equals = Leaf(token.EQUAL, "=") in split_leaves[0]
self.parenthesize_parent(new_node, parenth_before_equals)
node_to_split.replace(new_node)
return combined_prefix
