def replace_child(parent, node, replace_with): """Replace a node's child with another node while preserving formatting. Arguments: parent: (ast.AST) Parent node to replace a child of. node: (ast.AST) Child node to replace. replace_with: (ast.AST) New child node. """ # TODO(soupytwist): Don't refer to the formatting dict directly if hasattr(node, PASTA_DICT): setprop(replace_with, 'prefix', prop(node, 'prefix')) setprop(replace_with, 'suffix', prop(node, 'suffix')) for field in parent._fields: field_val = getattr(parent, field, None) if field_val == node: setattr(parent, field, replace_with) return elif isinstance(field_val, list): try: field_val[field_val.index(node)] = replace_with return except ValueError: pass raise errors.InvalidAstError('Node %r is not a child of %r' % (node, parent))
def split_import(sc, node, alias_to_remove): """Split an import node by moving the given imported alias into a new import. Arguments: sc: (scope.Scope) Scope computed on whole tree of the code being modified. node: (ast.Import|ast.ImportFrom) An import node to split. alias_to_remove: (ast.alias) The import alias node to remove. This must be a child of the given `node` argument. Raises: errors.InvalidAstError: if `node` is not appropriately contained in the tree represented by the scope `sc`. """ parent = sc.parent(node) parent_list = None for a in ('body', 'orelse', 'finalbody'): if hasattr(parent, a) and node in getattr(parent, a): parent_list = getattr(parent, a) break else: raise errors.InvalidAstError( 'Unable to find list containing import %r on ' 'parent node %r' % (node, parent)) idx = parent_list.index(node) new_import = copy.deepcopy(node) new_import.names = [alias_to_remove] node.names.remove(alias_to_remove) parent_list.insert(idx + 1, new_import) return new_import
def remove_child(parent, child): for _, field_value in ast.iter_fields(parent): if isinstance(field_value, list) and child in field_value: field_value.remove(child) return raise errors.InvalidAstError('Unable to find list containing child %r on ' 'parent node %r' % (child, parent))