def testRemoveAlias(self):
src = "from a import b, c"
tree = pasta.parse(src)
import_node = tree.body[0]
alias1 = import_node.names[0]
ast_utils.remove_child(import_node, alias1)
self.assertEqual(pasta.dump(tree), "from a import c")
def remove_import_alias_node(sc, node):
"""Remove an alias and if applicable remove their entire import.
Arguments:
sc: (scope.Scope) Scope computed on whole tree of the code being modified.
node: (ast.Import|ast.ImportFrom|ast.alias) The node to remove.
"""
import_node = sc.parent(node)
if len(import_node.names) == 1:
import_parent = sc.parent(import_node)
ast_utils.remove_child(import_parent, import_node)
else:
ast_utils.remove_child(import_node, node)
def remove_import(sc, alias_to_remove):
"""Remove an alias and if applicable remove their entire import.
Arguments:
sc: (scope.Scope) Scope computed on whole tree of the code being modified.
alias_to_remove: (list of ast.alias) The import alias nodes to remove.
"""
import_node = sc.parent(alias_to_remove)
if len(import_node.names) == 1:
import_parent = sc.parent(import_node)
ast_utils.remove_child(import_parent, import_node)
else:
ast_utils.remove_child(import_node, alias_to_remove)
def testRemoveFromBlock(self):
src = """\
if a:
print("foo!")
x = 1"""
tree = pasta.parse(src, py_ver)
if_block = tree.body[0]
print_stmt = if_block.body[0]
ast_utils.remove_child(if_block, print_stmt, py_ver=py_ver)
expected = """\
if a:
x = 1"""
self.assertEqual(pasta.dump(tree, py_ver), expected)
def testRemoveChildMethod(self):
src = """\
class C():
def f(x):
return x + 2
def g(x):
return x + 3"""
tree = pasta.parse(src, py_ver)
class_node = tree.body[0]
meth1_node = class_node.body[0]
ast_utils.remove_child(class_node, meth1_node, py_ver=py_ver)
result = pasta.dump(tree, py_ver)
expected = """\
class C():
def g(x):
return x + 3"""
self.assertEqual(result, expected)
def inline_name(t, name, py_ver=sys.version_info[:2]):
"""Inline a constant name into a module."""
sc = scope.analyze(t)
name_node = sc.names[name]
# The name must be a Name node (not a FunctionDef, etc.)
if not isinstance(name_node.definition, (ast27.Name, ast3.Name)):
raise InlineError('%r is not a constant; it has type %r' %
(name, type(name_node.definition)))
assign_node = sc.parent(name_node.definition)
if not isinstance(assign_node, (ast27.Assign, ast3.Assign)):
raise InlineError('%r is not declared in an assignment' % name)
value = assign_node.value
if not isinstance(sc.parent(assign_node), (ast27.Module, ast3.Module)):
raise InlineError('%r is not a top-level name' % name)
# If the name is written anywhere else in this module, it is not constant
for ref in name_node.reads:
if isinstance(getattr(ref, 'ctx', None), (ast27.Store, ast3.Store)):
raise InlineError('%r is not a constant' % name)
# Replace all reads of the name with a copy of its value
for ref in name_node.reads:
ast_utils.replace_child(sc.parent(ref), ref, copy.deepcopy(value))
# Remove the assignment to this name
if len(assign_node.targets) == 1:
ast_utils.remove_child(sc.parent(assign_node),
assign_node,
py_ver=py_ver)
else:
tgt_list = [
tgt for tgt in assign_node.targets
if not (isinstance(tgt, ast.Name) and tgt.id == name)
]
assign_node.targets = tgt_list
