def test_call_no_pos(self):
"""Tests that Call node traversal works without position information."""
src = 'f(a)'
t = pasta.parse(src, py_ver)
node = ast_utils.find_nodes_by_type(t, (ast27.Call, ast3.Call),
py_ver)[0]
node.keywords.append(
pasta.ast(py_ver).keyword(arg='b',
value=pasta.ast(py_ver).Num(n=0)))
self.assertEqual('f(a, b=0)', pasta.dump(t, py_ver))
def test_call_illegal_pos(self):
"""Tests that Call node traversal works even with illegal positions."""
src = 'f(a)'
t = pasta.parse(src, py_ver)
node = ast_utils.find_nodes_by_type(t, (ast27.Call, ast3.Call),
py_ver)[0]
node.keywords.append(
pasta.ast(py_ver).keyword(arg='b',
value=pasta.ast(py_ver).Num(n=0)))
# This position would put b=0 before a, so it should be ignored.
node.keywords[-1].value.lineno = 0
node.keywords[-1].value.col_offset = 0
self.assertEqual('f(a, b=0)', pasta.dump(t, py_ver))
def test_try_nested_imports(self):
source = textwrap.dedent("""\
try:
import aaa
except:
import bbb
finally:
import ccc
""")
tree = pasta.ast_parse(source, py_ver)
nodes = tree.body
node_aaa, node_bbb, node_ccc = ast_utils.find_nodes_by_type(
tree, pasta.ast(py_ver).alias, py_ver)
s = scope.analyze(tree, py_ver)
self.assertItemsEqual(s.names.keys(), {'aaa', 'bbb', 'ccc'})
self.assertItemsEqual(s.external_references.keys(), {'aaa', 'bbb', 'ccc'})
self.assertEqual(s.names['aaa'].definition, node_aaa)
self.assertEqual(s.names['bbb'].definition, node_bbb)
self.assertEqual(s.names['ccc'].definition, node_ccc)
for ref in {'aaa', 'bbb', 'ccc'}:
self.assertEqual(s.names[ref].reads, [],
'Expected no reads for %s' % ref)
class _TreeNormalizer(pasta.ast(py_ver).NodeTransformer):
"""Replaces all op nodes with unique instances."""
def visit(self, node):
if isinstance(node, _AST_OP_NODES):
return node.__class__()
return super(_TreeNormalizer, self).visit(node)
def test_autoindent(self):
src = textwrap.dedent("""\
def a():
b
c
""")
expected = textwrap.dedent("""\
def a():
b
new_node
""")
t = pasta.parse(src, py_ver)
# Repace the second node and make sure the indent level is corrected
t.body[0].body[1] = pasta.ast(py_ver).Expr(
pasta.ast(py_ver).Name(id='new_node'))
self.assertMultiLineEqual(expected, codegen.to_str(t, py_ver))
def remove_child(parent, child, py_ver=sys.version_info[:2]):
for _, field_value in pasta.ast(py_ver).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))
class FindNodeVisitor(pasta.ast(py_ver).NodeVisitor):
def __init__(self, condition):
self._condition = condition
self.results = []
def visit(self, node):
if self._condition(node):
self.results.append(node)
super(FindNodeVisitor, self).visit(node)
def make_safe_alias_node(alias_name, asname):
# Try to avoid name conflicts
new_alias = pasta.ast(py_ver).alias(name=alias_name, asname=asname)
imported_name = asname or alias_name
counter = 0
while imported_name in sc.names:
counter += 1
imported_name = new_alias.asname = '%s_%d' % (asname or alias_name,
counter)
return new_alias
def test_split_from_import(self):
src = 'from aaa import bbb, ccc, ddd\n'
t = pasta.ast_parse(src, py_ver)
import_node = t.body[0]
sc = scope.analyze(t, py_ver)
import_utils.split_import(sc, import_node, import_node.names[1])
self.assertEqual(2, len(t.body))
self.assertEqual(pasta.ast(py_ver).ImportFrom, type(t.body[1]))
self.assertEqual(t.body[0].module, 'aaa')
self.assertEqual(t.body[1].module, 'aaa')
self.assertEqual([alias.name for alias in t.body[0].names],
['bbb', 'ddd'])
def test_split_normal_import(self):
src = 'import aaa, bbb, ccc\n'
t = pasta.ast_parse(src, py_ver)
import_node = t.body[0]
sc = scope.analyze(t, py_ver)
import_utils.split_import(sc, import_node, import_node.names[1])
self.assertEqual(2, len(t.body))
self.assertEqual(pasta.ast(py_ver).Import, type(t.body[1]))
self.assertEqual([alias.name for alias in t.body[0].names],
['aaa', 'ccc'])
self.assertEqual([alias.name for alias in t.body[1].names],
['bbb'])
def test_functiondef_nested_imports(self):
source = textwrap.dedent("""\
def foo(bar):
import aaa
""")
tree = pasta.ast_parse(source, py_ver)
nodes = tree.body
node_aaa = ast_utils.find_nodes_by_type(
tree, pasta.ast(py_ver).alias, py_ver)[0]
s = scope.analyze(tree, py_ver)
self.assertItemsEqual(s.names.keys(), {'foo'})
self.assertItemsEqual(s.external_references.keys(), {'aaa'})
def test_remove_just_alias_import_from(self):
src = 'from m import a, b'
tree = pasta.ast_parse(src, py_ver)
sc = scope.analyze(tree, py_ver)
unused_b_node = tree.body[0].names[1]
import_utils.remove_import_alias_node(sc,
unused_b_node,
py_ver=py_ver)
self.assertEqual(len(tree.body), 1)
self.assertEqual(type(tree.body[0]), pasta.ast(py_ver).ImportFrom)
self.assertEqual(len(tree.body[0].names), 1)
self.assertEqual(tree.body[0].names[0].name, 'a')
def parse(src, py_ver=sys.version_info[:2]):
"""Replaces typed_ast.parse; ensures additional properties on the parsed tree.
This enforces the assumption that each node in the ast is unique.
"""
class _TreeNormalizer(pasta.ast(py_ver).NodeTransformer):
"""Replaces all op nodes with unique instances."""
def visit(self, node):
if isinstance(node, _AST_OP_NODES):
return node.__class__()
return super(_TreeNormalizer, self).visit(node)
tree=pasta.ast(py_ver).parse(sanitize_source(src))
_TreeNormalizer().visit(tree)
return tree
def test_lookup_scope(self):
src = textwrap.dedent("""\
import a
def b(c, d, e=1):
class F(d):
g = 1
return c
""")
t = pasta.ast_parse(src, py_ver)
import_node, func_node = t.body
class_node, return_node = func_node.body
sc = scope.analyze(t, py_ver)
import_node_scope = sc.lookup_scope(import_node)
self.assertIs(import_node_scope.node, t)
self.assertIs(import_node_scope, sc)
self.assertItemsEqual(import_node_scope.names, ['a', 'b'])
func_node_scope = sc.lookup_scope(func_node)
self.assertIs(func_node_scope.node, func_node)
self.assertIs(func_node_scope.parent_scope, sc)
self.assertItemsEqual(func_node_scope.names, ['c', 'd', 'e', 'F'])
class_node_scope = sc.lookup_scope(class_node)
self.assertIs(class_node_scope.node, class_node)
self.assertIs(class_node_scope.parent_scope, func_node_scope)
self.assertItemsEqual(class_node_scope.names, ['g'])
return_node_scope = sc.lookup_scope(return_node)
self.assertIs(return_node_scope.node, func_node)
self.assertIs(return_node_scope, func_node_scope)
self.assertItemsEqual(return_node_scope.names, ['c', 'd', 'e', 'F'])
self.assertIs(class_node_scope.lookup_scope(func_node), func_node_scope)
self.assertIsNone(sc.lookup_scope(pasta.ast(py_ver).Name(id='foo')))
class ScopeVisitor(pasta.ast(py_ver).NodeVisitor):
def __init__(self):
super(ScopeVisitor, self).__init__()
self._parent = None
self.root_scope = self.scope = RootScope(None)
def visit(self, node):
if node is None:
return
if self.root_scope.node is None:
self.root_scope.node = node
self.root_scope.set_parent(node, self._parent)
tmp = self._parent
self._parent = node
super(ScopeVisitor, self).visit(node)
self._parent = tmp
def visit_in_order(self, node, *attrs):
for attr in attrs:
val = getattr(node, attr, None)
if val is None:
continue
if isinstance(val, list):
for item in val:
self.visit(item)
elif isinstance(val, (ast27.AST, ast3.AST)):
self.visit(val)
def visit_Import(self, node):
for alias in node.names:
name_parts = alias.name.split('.')
if not alias.asname:
# If not aliased, define the top-level module of the import
cur_name = self.scope.define_name(name_parts[0], alias)
self.root_scope.add_external_reference(name_parts[0],
alias,
name_ref=cur_name)
# Define names of sub-modules imported
partial_name = name_parts[0]
for part in name_parts[1:]:
partial_name += '.' + part
cur_name = cur_name.lookup_name(part)
cur_name.define(alias)
self.root_scope.add_external_reference(
partial_name, alias, name_ref=cur_name)
else:
# If the imported name is aliased, define that name only
name = self.scope.define_name(alias.asname, alias)
# Define names of sub-modules imported
for i in range(1, len(name_parts)):
self.root_scope.add_external_reference(
'.'.join(name_parts[:i]), alias)
self.root_scope.add_external_reference(alias.name,
alias,
name_ref=name)
self.generic_visit(node)
def visit_ImportFrom(self, node):
if node.module:
name_parts = node.module.split('.')
for i in range(1, len(name_parts) + 1):
self.root_scope.add_external_reference(
'.'.join(name_parts[:i]), node)
for alias in node.names:
name = self.scope.define_name(alias.asname or alias.name,
alias)
if node.module:
self.root_scope.add_external_reference('.'.join(
(node.module, alias.name)),
alias,
name_ref=name)
# TODO: else? relative imports
self.generic_visit(node)
def visit_Name(self, node):
if isinstance(node.ctx,
(ast27.Store, ast3.Store, ast27.Param, ast3.Param)):
self.scope.define_name(node.id, node)
elif isinstance(node.ctx, (ast27.Load, ast3.Load)):
self.scope.lookup_name(node.id).add_reference(node)
self.root_scope.set_name_for_node(
node, self.scope.lookup_name(node.id))
self.generic_visit(node)
def visit_FunctionDef(self, node):
# Visit decorator list first to avoid declarations in args
self.visit_in_order(node, 'decorator_list')
if isinstance(self.root_scope.parent(node),
(ast27.ClassDef, ast3.ClassDef)):
pass # TODO: Support referencing methods by "self" where possible
else:
self.scope.define_name(node.name, node)
try:
self.scope = self.scope.create_scope(node)
self.visit_in_order(node, 'args', 'returns', 'body')
finally:
self.scope = self.scope.parent_scope
def visit_arguments(self, node):
self.visit_in_order(node, 'defaults', 'args')
if py_ver < (3, 0):
# In python 2.x, these names are not Name nodes. Define them explicitly
# to be able to find references in the function body.
for arg_attr_name in ('vararg', 'kwarg'):
arg_name = getattr(node, arg_attr_name, None)
if arg_name is not None:
self.scope.define_name(arg_name, node)
else:
# Visit defaults first to avoid declarations in args
self.visit_in_order(node, 'vararg', 'kwarg')
def visit_arg(self, node):
self.scope.define_name(node.arg, node)
# PEP 484 forward reference type annotations
if hasattr(node, 'annotation') and isinstance(
node.annotation, (ast27.Str, ast3.Str)):
name_parts = node.annotation.s.split('.')
# TODO: Fix this; the name may not be defined in the root scope.
name = self.root_scope.forward_define_name(
name_parts[0], node.annotation)
for part in name_parts[1:]:
name = name.lookup_name(part)
name.add_reference(node.annotation)
self.generic_visit(node)
def visit_ClassDef(self, node):
self.visit_in_order(node, 'decorator_list', 'bases')
self.scope.define_name(node.name, node)
try:
self.scope = self.scope.create_scope(node)
self.visit_in_order(node, 'body')
finally:
self.scope = self.scope.parent_scope
def visit_Attribute(self, node):
self.generic_visit(node)
node_value_name = self.root_scope.get_name_for_node(node.value)
if node_value_name:
node_name = node_value_name.lookup_name(node.attr)
self.root_scope.set_name_for_node(node, node_name)
node_name.add_reference(node)
def add_import(tree, name_to_import, asname=None, from_import=True, merge_from_imports=True,
py_ver=sys.version_info[:2]):
"""Adds an import to the module.
This function will try to ensure not to create duplicate imports. If name_to_import is
already imported, it will return the existing import. This is true even if asname is set
(asname will be ignored, and the existing name will be returned).
If the import would create a name that already exists in the scope given by tree, this
function will "import as", and append "_x" to the asname where x is the smallest positive
integer generating a unique name.
Arguments:
tree: (ast.Module) Module AST to modify.
name_to_import: (string) The absolute name to import.
asname: (string) The alias for the import ("import name_to_import as asname")
from_import: (boolean) If True, import the name using an ImportFrom node.
merge_from_imports: (boolean) If True, merge a newly inserted ImportFrom
node into an existing ImportFrom node, if applicable.
Returns:
The name (as a string) that can be used to reference the imported name. This
can be the fully-qualified name, the basename, or an alias name.
"""
sc = scope.analyze(tree, py_ver=py_ver)
# Don't add anything if it's already imported
if name_to_import in sc.external_references:
existing_ref = next((ref for ref in sc.external_references[name_to_import]
if ref.name_ref is not None), None)
if existing_ref:
return existing_ref.name_ref.id
import_node = None
added_name = None
def make_safe_alias_node(alias_name, asname):
# Try to avoid name conflicts
new_alias = pasta.ast(py_ver).alias(name=alias_name, asname=asname)
imported_name = asname or alias_name
counter = 0
while imported_name in sc.names:
counter += 1
imported_name = new_alias.asname = '%s_%d' % (asname or alias_name,
counter)
return new_alias
# Add an ImportFrom node if requested and possible
if from_import and '.' in name_to_import:
from_module, alias_name = name_to_import.rsplit('.', 1)
new_alias = make_safe_alias_node(alias_name, asname)
if merge_from_imports:
# Try to add to an existing ImportFrom from the same module
existing_from_import = next(
(node for node in tree.body
if isinstance(node, (ast27.ImportFrom, ast3.ImportFrom)) and
node.module == from_module and node.level == 0), None)
if existing_from_import:
existing_from_import.names.append(new_alias)
return new_alias.asname or new_alias.name
# Create a new node for this import
import_node = pasta.ast(py_ver).ImportFrom(
module=from_module, names=[new_alias], level=0)
# If not already created as an ImportFrom, create a normal Import node
if not import_node:
new_alias = make_safe_alias_node(alias_name=name_to_import, asname=asname)
import_node = pasta.ast(py_ver).Import(names=[new_alias])
# Insert the node at the top of the module and return the name in scope
tree.body.insert(1 if ast_utils.has_docstring(tree) else 0, import_node)
return new_alias.asname or new_alias.name
