def test_indent_levels_same_line(self):
src = 'if a: b; c\n'
t = pasta.parse(src)
if_node = t.body[0]
b, c = if_node.body
self.assertIsNone(fmt.get(b, 'indent_diff'))
self.assertIsNone(fmt.get(c, 'indent_diff'))
def test_indent_multiline_string_with_newline(self):
src = textwrap.dedent('''\
class A:
"""Doc\n
string."""
pass
''')
t = pasta.parse(src)
docstring, pass_stmt = t.body[0].body
self.assertEqual(' ', fmt.get(docstring, 'indent'))
self.assertEqual(' ', fmt.get(pass_stmt, 'indent'))
def test_fstring(self):
src = 'f"a {b} c d {e}"'
t = pasta.parse(src)
node = t.body[0].value
self.assertEqual(
fmt.get(node, 'content'),
'f"a {__pasta_fstring_val_0__} c d {__pasta_fstring_val_1__}"')
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_tabs_below_spaces_and_tab(self):
for num_spaces in range(1, 8):
t = pasta.parse(
textwrap.dedent('''\
if a:
{WS}{ONETAB}if b:
{ONETAB}{ONETAB}c
''').format(ONETAB='\t', WS=' ' * num_spaces))
node_c = t.body[0].body[0].body[0]
self.assertEqual(fmt.get(node_c, 'indent_diff'), '\t')
def test_indent_extra_newlines(self):
src = textwrap.dedent('''\
if a:
b
''')
t = pasta.parse(src)
if_node = t.body[0]
b = if_node.body[0]
self.assertEqual(' ', fmt.get(b, 'indent_diff'))
def test_scope_trailing_comma(self):
template = 'def foo(a, b{trailing_comma}): pass'
for trailing_comma in ('', ',', ' , '):
tree = pasta.parse(template.format(trailing_comma=trailing_comma))
self.assertEqual(
trailing_comma.lstrip(' ') + ')',
fmt.get(tree.body[0], 'args_suffix'))
template = 'class Foo(a, b{trailing_comma}): pass'
for trailing_comma in ('', ',', ' , '):
tree = pasta.parse(template.format(trailing_comma=trailing_comma))
self.assertEqual(
trailing_comma.lstrip(' ') + ')',
fmt.get(tree.body[0], 'bases_suffix'))
template = 'from mod import (a, b{trailing_comma})'
for trailing_comma in ('', ',', ' , '):
tree = pasta.parse(template.format(trailing_comma=trailing_comma))
self.assertEqual(trailing_comma + ')',
fmt.get(tree.body[0], 'names_suffix'))
def test_block_suffix(self):
src_tpl = textwrap.dedent('''\
{open_block}
pass #a
#b
#c
#d
#e
a
''')
test_cases = (
# first: attribute of the node with the last block
# second: code snippet to open a block
('body', 'def x():'),
('body', 'class X:'),
('body', 'if x:'),
('orelse', 'if x:\n y\nelse:'),
('body', 'if x:\n y\nelif y:'),
('body', 'while x:'),
('orelse', 'while x:\n y\nelse:'),
('finalbody', 'try:\n x\nfinally:'),
('body', 'try:\n x\nexcept:'),
('orelse', 'try:\n x\nexcept:\n y\nelse:'),
('body', 'with x:'),
('body', 'with x, y:'),
('body', 'with x:\n with y:'),
('body', 'for x in y:'),
)
def is_node_for_suffix(node, children_attr):
# Return True if this node contains the 'pass' statement
val = getattr(node, children_attr, None)
return isinstance(val, list) and type(val[0]) == ast.Pass
for children_attr, open_block in test_cases:
src = src_tpl.format(open_block=open_block)
t = pasta.parse(src)
node_finder = ast_utils.FindNodeVisitor(
lambda node: is_node_for_suffix(node, children_attr))
node_finder.visit(t)
node = node_finder.results[0]
expected = ' #b\n #c\n\n #d\n'
actual = str(fmt.get(node, 'block_suffix_%s' % children_attr))
self.assertMultiLineEqual(
expected, actual,
'Incorrect suffix for code:\n%s\nNode: %s (line %d)\nDiff:\n%s'
% (src, node, node.lineno, '\n'.join(
_get_diff(actual, expected))))
self.assertMultiLineEqual(src, pasta.dump(t))
def test(self):
with open(input_file, 'r') as handle:
src = handle.read()
t = ast_utils.parse(src)
annotator = annotate.AstAnnotator(src)
annotator.visit(t)
def escape(s):
return '' if s is None else s.replace('\n', '\\n')
result = '\n'.join(
"{0:12} {1:20} \tprefix=|{2}|\tsuffix=|{3}|\tindent=|{4}|".
format(
str((getattr(n, 'lineno', -1),
getattr(n, 'col_offset', -1))),
type(n).__name__ + ' ' +
_get_node_identifier(n), escape(fmt.get(n, 'prefix')),
escape(fmt.get(n, 'suffix')),
escape(fmt.get(n, 'indent')))
for n in ast.walk(t)) + '\n'
# If specified, write the golden data instead of checking it
if getattr(self, 'generate_goldens', False):
if not os.path.isdir(os.path.dirname(golden_file)):
os.makedirs(os.path.dirname(golden_file))
with open(golden_file, 'w') as f:
f.write(result)
print('Wrote: ' + golden_file)
return
try:
with io.open(golden_file, 'r', encoding='UTF-8') as f:
golden = f.read()
except IOError:
self.fail('Missing golden data.')
self.assertMultiLineEqual(golden, result)
def test_indent_depths(self):
template = 'if a:\n{first}if b:\n{first}{second}foo()\n'
indents = (' ', ' ' * 2, ' ' * 4, ' ' * 8, '\t', '\t' * 2)
for first, second in itertools.product(indents, indents):
src = template.format(first=first, second=second)
t = pasta.parse(src)
outer_if_node = t.body[0]
inner_if_node = outer_if_node.body[0]
call_node = inner_if_node.body[0]
self.assertEqual('', fmt.get(outer_if_node, 'indent'))
self.assertEqual('', fmt.get(outer_if_node, 'indent_diff'))
self.assertEqual(first, fmt.get(inner_if_node, 'indent'))
self.assertEqual(first, fmt.get(inner_if_node, 'indent_diff'))
self.assertEqual(first + second, fmt.get(call_node, 'indent'))
self.assertEqual(second, fmt.get(call_node, 'indent_diff'))
def test_no_block_suffix_for_single_line_statement(self):
src = 'if x: return y\n #a\n#b\n'
t = pasta.parse(src)
self.assertIsNone(fmt.get(t.body[0], 'block_suffix_body'))
def test_module_suffix(self):
src = 'foo\n#bar\n\n#baz\n'
t = pasta.parse(src)
self.assertEqual(src[src.index('#bar'):], fmt.get(t, 'suffix'))
def test_fstring_escaping(self):
src = 'f"a {{{b} {{c}}"'
t = pasta.parse(src)
node = t.body[0].value
self.assertEqual(fmt.get(node, 'content'),
'f"a {{{__pasta_fstring_val_0__} {{c}}"')
def test_statement_prefix_suffix(self):
src = 'a\n\ndef foo():\n return bar\n\n\nb\n'
t = pasta.parse(src)
self.assertEqual('\n', fmt.get(t.body[1], 'prefix'))
self.assertEqual('', fmt.get(t.body[1], 'suffix'))
def test_expression_prefix_suffix(self):
src = 'a\n\nfoo\n\n\nb\n'
t = pasta.parse(src)
self.assertEqual('\n', fmt.get(t.body[1], 'prefix'))
self.assertEqual('\n', fmt.get(t.body[1], 'suffix'))
