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))
