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)