def add_provide_context_to_python_operator(node: LN, capture: Capture, filename: Filename) -> None:
fn_args = capture['function_arguments'][0]
fn_args.append_child(Comma())
provide_context_arg = KeywordArg(Name('provide_context'), Name('True'))
provide_context_arg.prefix = fn_args.children[0].prefix
fn_args.append_child(provide_context_arg)
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 explode(self, is_def: bool, prefix: str = "") -> List[LN]:
result: List[LN] = []
prefix = self.prefix if self.prefix else prefix
if is_def:
if self.star:
self.star.prefix = prefix
result.append(self.star)
prefix = ""
if self.annotation:
result.append(
Node(
SYMBOL.tname,
[
Name(self.name, prefix=prefix),
Leaf(TOKEN.COLON, ":", prefix=""),
Name(self.annotation, prefix=" "),
],
prefix=prefix,
))
else:
result.append(Name(self.name, prefix=prefix))
if self.value:
space = " " if self.annotation else ""
result.append(Leaf(TOKEN.EQUAL, "=", prefix=space))
result.append(self.value)
else:
if self.star:
if self.star.type == TOKEN.STAR:
node = Node(SYMBOL.star_expr, [self.star], prefix=prefix)
elif self.star.type == TOKEN.DOUBLESTAR:
node = Node(SYMBOL.argument, [self.star], prefix=prefix)
if self.value:
self.value.prefix = ""
node.append_child(self.value)
result.append(node)
return result
if self.name:
self.value.prefix = ""
result.append(
Node(
SYMBOL.argument,
[
Name(self.name, prefix=prefix),
Leaf(TOKEN.EQUAL, value="=", prefix=""),
self.value,
],
prefix=prefix,
))
else:
self.value.prefix = prefix
result.append(self.value)
return result
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)]
def add_provide_context_to_python_operator(node: LN, capture: Capture, filename: Filename) -> None:
fn_args = capture['function_arguments'][0]
if len(fn_args.children) > 0 and (not isinstance(fn_args.children[-1], Leaf)
or fn_args.children[-1].type != token.COMMA):
fn_args.append_child(Comma())
provide_context_arg = KeywordArg(Name('provide_context'), Name('True'))
provide_context_arg.prefix = fn_args.children[0].prefix
fn_args.append_child(provide_context_arg)
def process_root(node: LN, capture: Capture,
filename: Filename) -> Optional[LN]:
print(filename)
# Find any imports close to the root level
imports = [
x for x in get_children(
node, python_symbols.import_from, recursive=True, recurse_depth=1)
if x.children[1].type == token.NAME
and x.children[1].value == "__future__"
]
if len(imports) <= 1 and FLOAT_MODE:
# Do nothing if one or no imports and only in float-mode
return
elif len(imports) == 0:
# We're in update mode - add missing __future__ imports
import_insert_point = find_future_import_insert_point(node)
insert = driver.parse_string(
"from __future__ import absolute_import, division, print_function\n"
)
node.children.insert(import_insert_point, insert)
return
# Extract the list of __future__ imports from all of the import statements
future_names = set()
for imp in imports:
if imp.children[3].type == token.NAME:
future_names.add(imp.children[3].value)
elif imp.children[3].type == python_symbols.import_as_names:
future_names.update(
{x.value
for x in get_children(imp.children[3], token.NAME)})
# If we're not purely floating, update the list of names
if not FLOAT_MODE:
# If present, make all the basics present
if future_names:
future_names |= {"absolute_import", "division", "print_function"}
# Remove anything already mandatory
future_names -= {"generators", "nested_scopes", "with_statement"}
# Update the first import instance with all the names
if len(future_names) == 1:
imports[0].children[3] = Name(list(future_names)[0])
else:
# Make the first one a multi-import
commad = join_comma(
[Name(x, prefix=" ") for x in sorted(future_names)])
names = Node(python_symbols.import_as_names, commad)
imports[0].children[3] = names
# Remove the other imports
for imp in imports[1:]:
assert imp.parent.type == python_symbols.simple_stmt
imp.parent.parent.children.remove(imp.parent)
def transform(self, node, results):
if self.should_skip(node):
return
touch_import("six.moves", "input", node)
name = results["name"]
if name.value == "raw_input":
name.replace(Name("input", prefix=name.prefix))
else:
new_node = node.clone()
new_node.prefix = ""
return Call(Name("eval"), [new_node], prefix=node.prefix)
def transform(self, node, results):
if self.should_skip(node):
return
touch_import(u'six.moves', u'input', node)
name = results['name']
if name.value == 'raw_input':
name.replace(Name('input', prefix=name.prefix))
else:
new_node = node.clone()
new_node.prefix = ''
return Call(Name('eval'), [new_node], prefix=node.prefix)
def transform_import(self, node, results):
"""Transform for the basic import case. Replaces the old
import name with a comma separated list of its
replacements.
"""
import_mod = results.get("module")
pref = import_mod.prefix
names = []
# create a Node list of the replacement modules
for name in MAPPING[import_mod.value][:-1]:
names.extend([Name(name[0], prefix=pref), Comma()])
names.append(Name(MAPPING[import_mod.value][-1][0], prefix=pref))
import_mod.replace(names)
def on_delete_modifier(node, capture, filename) -> None:
if not capture:
return
args: Node = capture.get("function_arguments")[0]
has_on_delete = False
for node in args.children:
if isinstance(node, Node) and node.type == 261:
if node.children[0].value == "on_delete":
has_on_delete = True
if not has_on_delete:
if args.children[-1].value != ",":
args.append_child(Comma())
args.append_child(
KeywordArg(Name("on_delete", prefix=" "), Name("models.CASCADE"))
)
def transform(self, node, results):
if _literal_re.match(node.value):
fixer_util.touch_import(None, "six", node)
new = node.clone()
new.value = new.value[1:]
new.prefix = ""
node.replace(Call(Name("six.u", prefix=node.prefix), [new]))
def replace_decorators(node, capture, filename):
"""
Replaces usage of ``@tornado.<etc>`` with ``@salt.ext.tornado.<etc>``
"""
indent = find_indentation(node)
decorator = _get_decorator(node)
decorator.remove()
decorated = Node(
SYMBOL.decorated,
[
Node(
SYMBOL.decorator,
[
Leaf(TOKEN.AT, "@"),
Name("salt.ext.{}".format(get_decorator_name(decorator))),
Leaf(TOKEN.NEWLINE, "\n"),
],
)
],
prefix=decorator.prefix,
)
node.replace(decorated)
decorated.append_child(node)
if indent is not None:
node.prefix = indent
else:
node.prefix = ""
def _add_arg_name(argument_node):
nonlocal encounter_kwarg
nonlocal idx
if args_list is None:
return
if encounter_kwarg:
return
if idx >= len(args_list):
msg = 'args_list: "{}" is shorter than positional arguments.'.format(
args_list)
log_error(filename, argument_node.get_lineno(), msg)
return
if len(argument_node.children) >= 3:
encounter_kwarg = True
msg = 'args_list: "{}" is longer than positional arguments, redundant arguments will be skipped.'.format(
args_list)
log_info(filename, argument_node.get_lineno(), msg)
return
key = args_list[idx]
argument_node.insert_child(0, Leaf(token.EQUAL, "="))
argument_node.insert_child(0, Name(key))
argument_node.children[0].prefix = argument_node.children[2].prefix
argument_node.children[2].prefix = ""
idx += 1
msg = 'add argument name "{}" for {}-th argument.'.format(key, idx)
log_debug(filename, argument_node.get_lineno(), msg)
def transform(self, node, results):
if _literal_re.match(node.value):
touch_import(None, u'six', node)
new = node.clone()
new.value = new.value[1:]
new.prefix = ''
node.replace(Call(Name(u'six.u', prefix=node.prefix), [new]))
def add_argument_transform(
node: Node, capture: Capture, filename: Filename
) -> None:
if "function_def" not in capture and "function_call" not in capture:
return
spec = FunctionSpec.build(node, capture)
done = False
value_leaf = Name(value)
if spec.is_def:
new_arg = FunctionArgument(
name,
value_leaf if keyword else None,
cast(str, type_annotation) if type_annotation != SENTINEL else "",
)
for index, argument in enumerate(spec.arguments):
if after == argument.name:
spec.arguments.insert(index + 1, new_arg)
done = True
break
if (
after == START
or (positional and (argument.value or argument.star))
or (
keyword
and argument.star
and argument.star.type == TOKEN.DOUBLESTAR
)
):
spec.arguments.insert(index, new_arg)
done = True
break
if not done:
spec.arguments.append(new_arg)
elif positional:
new_arg = FunctionArgument(value=value_leaf)
for index, argument in enumerate(spec.arguments):
if argument.star and argument.star.type == TOKEN.STAR:
log.debug(f"noping out due to *{argument.name}")
done = True
break
if index == stop_at:
spec.arguments.insert(index + 1, new_arg)
done = True
break
if after == START or argument.name or argument.star:
spec.arguments.insert(index, new_arg)
done = True
break
if not done:
spec.arguments.append(new_arg)
spec.explode()
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 _Call(self, name, args=None, prefix=None):
"""Help the next test"""
children = []
if isinstance(args, list):
for arg in args:
children.append(arg)
children.append(Comma())
children.pop()
return Call(Name(name), children, prefix)
def add_argument(filename, trailer_node, key, value):
"""
add "key=value" to arglist in trailer_node,
if arglist already contains key, reassign key to value
"""
if trailer_node.type != python_symbols.trailer and len(
trailer_node.children) != 3:
log_warning(
filename, trailer_node.get_lineno(),
"node type is not trailer or len(children) != 3. you may need to call norm_arglist first."
)
return
arglist_node = trailer_node.children[1]
if arglist_node.type != python_symbols.arglist:
log_warning(filename, trailer_node.get_lineno(),
"trailer_node.children[1] is not arglist.")
return
found_key = False
for node in arglist_node.children:
if node.type != python_symbols.argument:
continue
_key_node = node.children[0]
if _key_node.value == key:
found_key = True
_value_node = node.children[2]
if _value_node.value != value:
_value_node_copy = _value_node.clone()
_value_node_copy.type = token.NAME
_value_node_copy.value = value
_value_node.replace(_value_node_copy)
log_warning(
filename, arglist_node.get_lineno(),
'argument "{}" is reassigned to "{}"'.format(key, value))
break
if not found_key:
key_node = Name(key)
value_node = Name(value)
if arglist_node.children:
arglist_node.append_child(Comma())
key_node.prefix = " "
arg_node = KeywordArg(key_node, value_node)
arglist_node.append_child(arg_node)
log_warning(filename, arglist_node.get_lineno(),
'add argument "{}={}"'.format(key, value))
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 transform(self, node, results):
exc = results["exc"].clone()
val = results["val"].clone()
tb = results["tb"].clone()
exc.prefix = u""
val.prefix = tb.prefix = u" "
touch_import(None, u'six', node)
return Call(Name(u"six.reraise"),
[exc, Comma(), val, Comma(), tb],
prefix=node.prefix)
def modify_attr(node: LN, capture: Capture,
filename: Filename) -> Optional[LN]:
node.replace(
Call(
Name("getattr"),
args=[
capture["obj"].clone(),
Comma(),
Space(),
String('"' + capture["attr"].value + '"'),
],
))
def rewrite_print(node, capture, filename):
"""Given a print node, rewrite to a logger.debug"""
params = capture["function_parameters"]
# Args is either a Leaf or an arglist Node here
arglist = params.children[1]
# Extract the arguments from this list
if isinstance(arglist, Node) and arglist.type == syms.arglist:
args = [x for x in arglist.children if not x.type == token.COMMA]
# Remove kwargs for now
non_kwargs = [x for x in args if not x.type == syms.argument]
# Read out things like sep here
sep = " "
if not len(non_kwargs) == len(args):
first_leaf = next(node.leaves())
print(
"Warning: {}:{} Not handling keyword argument loggers".format(
filename, first_leaf.lineno))
return None
# If more than one child, we need to construct a new joiner string
if len(non_kwargs) > 1:
# Instead of placing a new, modify the first if it's a string
if arglist.children[0].type == token.STRING:
value = arglist.children[0].value
last_char = value[-1]
new_end = sep + sep.join(["%s"] * (len(non_kwargs) - 1))
arglist.children[0].value = value[:-1] + new_end + last_char
else:
arglist.insert_child(
0, String('"' + sep.join(["%s"] * len(non_kwargs)) + '"'))
arglist.insert_child(1, Comma())
arglist.children[2].prefix = " " + arglist.children[2].prefix
# Use the possibly rewritten parameters in the new call
new_node = Attr(Name("logger"), Name("debug"))
new_node[0].prefix = node.children[0].prefix
node.children = new_node + [params]
def convert_regex_to_path_modifier(node: Node, capture: Capture,
filename: Filename) -> None:
# Replace the import
if is_import(node):
name_leafs = [
Name("path", prefix=" "),
Comma(),
Name("re_path", prefix=" "),
]
node.replace([FromImport("django.url", name_leafs=name_leafs)])
# And function calls from url to path, re_path
if capture and "function_arguments" in capture:
function_node: Node = next(node.leaves())
args = capture.get("function_arguments")
regex_leaf: Leaf = next(args[0].leaves())
converted = update_regex_to_path(regex_leaf.value)
if converted == regex_leaf.value:
function_node.replace(Name("re_path", prefix=function_node.prefix))
else:
function_node.replace(Name("path", prefix=function_node.prefix))
regex_leaf.value = update_regex_to_path(regex_leaf.value)
def rename_transform(node: LN, capture: Capture,
filename: Filename) -> None:
log.debug(f"{filename} [{list(capture)}]: {node}")
for _key, value in capture.items():
log.debug(f"{_key}: {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":
parts = zip(
dotted_parts(old_name),
dotted_parts(new_name),
value.children,
)
for old, new, leaf in parts:
if old != leaf.value:
break
if old != new:
leaf.replace(Name(new, prefix=leaf.prefix))
elif type_repr(value.type) == "power":
pows = zip(dotted_parts(old_name),
dotted_parts(new_name))
it = iter(value.children)
leaf = next(it)
for old, new in pows:
if old == ".":
leaf = leaf.children[1]
continue
if old != leaf.value:
break
if old == leaf.value and old != new:
leaf.replace(Name(new, prefix=leaf.prefix))
leaf = next(it)
def modify_dict_literal(node: LN, capture: Capture,
filename: Filename) -> Optional[LN]:
toks = iter(capture.get("body"))
items = []
prefix = ""
while True:
try:
tok = next(toks)
if tok.type == TOKEN.DOUBLESTAR:
body = next(toks).clone()
body.prefix = prefix + tok.prefix + body.prefix
items.append(body)
else:
colon = next(toks)
value = next(toks).clone()
value.prefix = colon.prefix + value.prefix
if items and isinstance(items[-1], list):
items[-1].append(TupleNode(tok, value))
else:
items.append([TupleNode(tok, value)])
comma = next(toks)
prefix = comma.prefix
except StopIteration:
break
listitems = []
for item in items:
if listitems:
listitems.extend([Space(), Plus(), Space()])
if isinstance(item, list):
listitems.append(ListNode(*item))
else:
call = Node(SYMBOL.test, [*Attr(item, Name("items")), ArgList([])])
listitems.append(call)
args = listitems
if len(listitems) > 1:
args = [Node(SYMBOL.arith_expr, args)]
args.append(String(node.children[-1].prefix))
return Call(Name("dict"), args, prefix=node.prefix)
def modify_argument_transform(node: Node, capture: Capture,
filename: Filename) -> None:
spec = FunctionSpec.build(node, capture)
for argument in spec.arguments:
if argument.name == name:
if new_name is not SENTINEL:
argument.name = str(new_name)
if spec.is_def and type_annotation is not SENTINEL:
argument.annotation = str(type_annotation)
if spec.is_def and default_value is not SENTINEL:
argument.value = Name(default_value, prefix=" ")
spec.explode()
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_dot(self, node, results):
"""Transform for calls to module members in code."""
module_dot = results.get("bare_with_attr")
member = results.get("member")
new_name = None
if isinstance(member, list):
member = member[0]
for change in MAPPING[module_dot.value]:
if member.value in change[1]:
new_name = change[0]
break
if new_name:
module_dot.replace(Name(new_name, prefix=module_dot.prefix))
else:
self.cannot_convert(node, "This is an invalid module element")
def __call__(self, node: LN, capture: Capture, filename: Filename) -> None:
old_node = capture['module_name']
new_node = Node(
type=syms.dotted_as_name,
children=[
build_new_name_node(
old_node=old_node,
new_name=self.new_name,
attach=False,
),
Name('as', prefix=' '),
old_node.clone(),
],
)
old_node.replace(new_node)
