def test_split_nested_imports(self):
test_cases = (
'def foo():\n {import_stmt}\n',
'class Foo(object):\n {import_stmt}\n',
'if foo:\n {import_stmt}\nelse:\n pass\n',
'if foo:\n pass\nelse:\n {import_stmt}\n',
'if foo:\n pass\nelif bar:\n {import_stmt}\n',
'try:\n {import_stmt}\nexcept:\n pass\n',
'try:\n pass\nexcept:\n {import_stmt}\n',
'try:\n pass\nfinally:\n {import_stmt}\n',
'for i in foo:\n {import_stmt}\n',
'for i in foo:\n pass\nelse:\n {import_stmt}\n',
'while foo:\n {import_stmt}\n',
)
for template in test_cases:
try:
src = template.format(import_stmt='import aaa, bbb, ccc')
t = ast.parse(src)
sc = scope.analyze(t)
import_node = ast_utils.find_nodes_by_type(t, ast.Import)[0]
import_utils.split_import(sc, import_node, import_node.names[1])
split_import_nodes = ast_utils.find_nodes_by_type(t, ast.Import)
self.assertEqual(1, len(t.body))
self.assertEqual(2, len(split_import_nodes))
self.assertEqual([alias.name for alias in split_import_nodes[0].names],
['aaa', 'ccc'])
self.assertEqual([alias.name for alias in split_import_nodes[1].names],
['bbb'])
except:
self.fail('Failed while executing case:\n%s\nCaused by:\n%s' %
(src, traceback.format_exc()))
def test_try_nested_imports(self):
source = textwrap.dedent("""\
try:
import aaa
except:
import bbb
finally:
import ccc
""")
tree = ast.parse(source)
nodes = tree.body
node_aaa, node_bbb, node_ccc = ast_utils.find_nodes_by_type(
tree, ast.alias)
s = scope.analyze(tree)
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)
def test_indent_levels(self):
src = textwrap.dedent('''\
foo('begin')
if a:
foo('a1')
if b:
foo('b1')
if c:
foo('c1')
foo('b2')
foo('a2')
foo('end')
''')
t = pasta.parse(src)
call_nodes = ast_utils.find_nodes_by_type(t, (ast.Call, ))
call_nodes.sort(key=lambda node: node.lineno)
begin, a1, b1, c1, b2, a2, end = call_nodes
self.assertEqual('', fmt.get(begin, 'indent'))
self.assertEqual(' ', fmt.get(a1, 'indent'))
self.assertEqual(' ', fmt.get(b1, 'indent'))
self.assertEqual(' ', fmt.get(c1, 'indent'))
self.assertEqual(' ', fmt.get(b2, 'indent'))
self.assertEqual(' ', fmt.get(a2, 'indent'))
self.assertEqual('', fmt.get(end, 'indent'))
def test_if_nested_imports(self):
source = textwrap.dedent("""\
if a:
import aaa
elif b:
import bbb
else:
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', 'a', 'b'})
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)
self.assertIsNone(s.names['a'].definition)
self.assertIsNone(s.names['b'].definition)
for ref in {'aaa', 'bbb', 'ccc'}:
self.assertEqual(s.names[ref].reads, [],
'Expected no reads for %s' % ref)
def test_call_no_pos(self):
"""Tests that Call node traversal works without position information."""
src = 'f(a)'
t = pasta.parse(src)
node = ast_utils.find_nodes_by_type(t, (ast.Call, ))[0]
node.keywords.append(ast.keyword(arg='b', value=ast.Num(n=0)))
self.assertEqual('f(a, b=0)', pasta.dump(t))
def test_call_illegal_pos(self):
"""Tests that Call node traversal works even with illegal positions."""
src = 'f(a)'
t = pasta.parse(src)
node = ast_utils.find_nodes_by_type(t, (ast.Call, ))[0]
node.keywords.append(ast.keyword(arg='b', value=ast.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))
def test_functiondef_nested_imports(self):
source = textwrap.dedent("""\
def foo(bar):
import aaa
""")
tree = ast.parse(source)
nodes = tree.body
node_aaa = ast_utils.find_nodes_by_type(tree, ast.alias)[0]
s = scope.analyze(tree)
self.assertItemsEqual(s.names.keys(), {'foo'})
self.assertItemsEqual(s.external_references.keys(), {'aaa'})
def test_indent_levels(self):
src = textwrap.dedent("""\
foo('begin',
bar('a1'))
bar('x')
""")
t = pasta.parse(src)
call_nodes = ast_utils.find_nodes_by_type(t, _ast.Call)
call_nodes.sort(key=lambda node: node.lineno)
foo, bar_arg, bar = call_nodes
self.assertEqual(1, fmt.get(foo, 'start_line'))
self.assertEqual(0, fmt.get(foo, 'start_col'))
self.assertEqual(2, fmt.get(foo, 'end_line'))
self.assertEqual(13, fmt.get(foo, 'end_col'))
self.assertEqual(2, fmt.get(bar_arg, 'start_line'))
self.assertEqual(2, fmt.get(bar_arg, 'start_col'))
self.assertEqual(2, fmt.get(bar_arg, 'end_line'))
self.assertEqual(12, fmt.get(bar_arg, 'end_col'))
self.assertEqual(3, fmt.get(bar, 'start_line'))
self.assertEqual(0, fmt.get(bar, 'start_col'))
self.assertEqual(3, fmt.get(bar, 'end_line'))
self.assertEqual(9, fmt.get(bar, 'end_col'))
