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(ast_utils.prop(b, 'indent_diff'))
self.assertIsNone(ast_utils.prop(c, 'indent_diff'))
def attr(self, node, attr_name, attr_vals, deps=None, default=None):
"""Add the formatted data stored for a given attribute on this node.
If any of the dependent attributes of the node have changed since it was
annotated, then the stored formatted data for this attr_name is no longer
valid, and we must use the default instead.
Arguments:
node: (ast.AST) An AST node to retrieve formatting information from.
attr_name: (string) Name to load the formatting information from.
attr_vals: (list of functions/strings) Unused here.
deps: (optional, set of strings) Attributes of the node which the stored
formatting data depends on.
default: (string) Default formatted data for this attribute.
"""
del attr_vals
if not hasattr(node, '_printer_info') or node._printer_info[attr_name]:
return
node._printer_info[attr_name] = True
if (deps and any(
getattr(node, dep, None) != ast_utils.prop(
node, dep + '__src') for dep in deps)):
self.code += default or ''
else:
val = ast_utils.prop(node, attr_name)
self.code += val if val is not None else (default or '')
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}|".format(
str((getattr(n, 'lineno', -1), getattr(n, 'col_offset', -1))),
type(n).__name__ + ' ' + _get_node_identifier(n),
escape(ast_utils.prop(n, 'prefix')),
escape(ast_utils.prop(n, 'suffix')))
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 open(golden_file, 'r') as f:
golden = f.read()
except IOError:
self.fail('Missing golden data.')
self.assertMultiLineEqual(golden, result)
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('', ast_utils.prop(begin, 'indent'))
self.assertEqual(' ', ast_utils.prop(a1, 'indent'))
self.assertEqual(' ', ast_utils.prop(b1, 'indent'))
self.assertEqual(' ', ast_utils.prop(c1, 'indent'))
self.assertEqual(' ', ast_utils.prop(b2, 'indent'))
self.assertEqual(' ', ast_utils.prop(a2, 'indent'))
self.assertEqual('', ast_utils.prop(end, 'indent'))
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(' ') + ')',
ast_utils.prop(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(' ') + ')',
ast_utils.prop(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 + ')',
ast_utils.prop(tree.body[0], 'names_suffix'))
def optional_token(self,
node,
attr_name,
token_val,
allow_whitespace_prefix=False):
del token_val, allow_whitespace_prefix
if not hasattr(node, ast_utils.PASTA_DICT):
return
self.code += ast_utils.prop(node, attr_name) or ''
def close_scope(self,
node,
prefix_attr='prefix',
suffix_attr='suffix',
trailing_comma=False):
"""Close a parenthesized scope on the given node, if one is open."""
# Ensures the prefix + suffix are not None
if ast_utils.prop(node, prefix_attr) is None:
ast_utils.setprop(node, prefix_attr, '')
if ast_utils.prop(node, suffix_attr) is None:
ast_utils.setprop(node, suffix_attr, '')
if not self._parens or node not in self._scope_stack[-1]:
return
symbols = {')'}
if trailing_comma:
symbols.add(',')
parsed_to_i = self._i
parsed_to_loc = prev_loc = self._loc
result = ''
for tok in self.takewhile(
lambda t: t.type in FORMATTING_TOKENS or t.src in symbols):
# Stores all the code up to and including this token
result += self._space_between(prev_loc, tok) + tok.src
if tok.src == ')':
# Close out the open scope
self._scope_stack.pop()
ast_utils.prependprop(node, prefix_attr, self._parens.pop())
ast_utils.appendprop(node, suffix_attr, result)
result = ''
parsed_to_i = self._i
parsed_to_loc = tok.end
if not self._parens or node not in self._scope_stack[-1]:
break
prev_loc = tok.end
# Reset back to the last place where we parsed anything
self._i = parsed_to_i
self._loc = parsed_to_loc
def wrapped(self, node, *args, **kwargs):
self.prefix(node)
f(self, node, *args, **kwargs)
if hasattr(self, 'block_suffix'):
last_child = ast_utils.get_last_child(node)
# Workaround for ast.Module which does not have a lineno
if last_child and last_child.lineno != getattr(node, 'lineno', 0):
indent = (ast_utils.prop(last_child, 'prefix')
or '\n').splitlines()[-1]
self.block_suffix(node, indent)
else:
self.suffix(node, comment=True)
def test_block_suffix(self):
src_tpl = textwrap.dedent('''\
{open_block}
pass #a
#b
#c
#d
#e
a
''')
test_cases = (
'def x():',
'class X:',
'if x:',
'if x:\n y\nelse:',
'if x:\n y\nelif y:',
'while x:',
'while x:\n y\nelse:',
'try:\n x\nfinally:',
'try:\n x\nexcept:',
'try:\n x\nexcept:\n y\nelse:',
'with x:',
'with x, y:',
'with x:\n with y:',
'for x in y:',
)
def is_node_for_suffix(node):
# Return True if this node contains the 'pass' statement
for attr in dir(node):
attr_val = getattr(node, attr)
if (isinstance(attr_val, list) and any(
isinstance(child, ast.Pass) for child in attr_val)):
return True
return False
node_finder = ast_utils.FindNodeVisitor(is_node_for_suffix)
for open_block in test_cases:
src = src_tpl.format(open_block=open_block)
t = pasta.parse(src)
node_finder.results = []
node_finder.visit(t)
node = node_finder.results[0]
expected = ' #b\n #c\n\n #d\n'
actual = ast_utils.prop(node, 'suffix')
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))))
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(
ast_utils.prop(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))))
def visit_If(self, node):
tok = 'elif' if ast_utils.prop(node, 'is_elif') else 'if'
self.attr(node, 'open_if', [tok, self.ws], default=tok + ' ')
self.visit(node.test)
self.attr(node,
'open_block', [self.ws, ':', self.ws_oneline],
default=':\n')
for stmt in node.body:
self.visit(stmt)
if node.orelse:
if (len(node.orelse) == 1 and isinstance(node.orelse[0], ast.If)
and self.check_is_elif(node.orelse[0])):
ast_utils.setprop(node.orelse[0], 'is_elif', True)
self.visit(node.orelse[0])
else:
self.attr(node, 'elseprefix', [self.ws])
self.token('else')
self.attr(node,
'open_else', [self.ws, ':', self.ws_oneline],
default=':\n')
for stmt in node.orelse:
self.visit(stmt)
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('', ast_utils.prop(outer_if_node, 'indent'))
self.assertEqual('', ast_utils.prop(outer_if_node, 'indent_diff'))
self.assertEqual(first, ast_utils.prop(inner_if_node, 'indent'))
self.assertEqual(first, ast_utils.prop(inner_if_node,
'indent_diff'))
self.assertEqual(first + second,
ast_utils.prop(call_node, 'indent'))
self.assertEqual(second, ast_utils.prop(call_node, 'indent_diff'))
def test_expression_prefix_suffix(self):
src = 'a\n\nfoo\n\n\nb\n'
t = pasta.parse(src)
self.assertEqual('\n', ast_utils.prop(t.body[1], 'prefix'))
self.assertEqual('\n', ast_utils.prop(t.body[1], 'suffix'))
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', ast_utils.prop(t.body[1], 'prefix'))
self.assertEqual('', ast_utils.prop(t.body[1], 'suffix'))
def optional_token(self, node, attr_name, token_val):
del token_val
if not hasattr(node, ast_utils.PASTA_DICT):
return
self.code += ast_utils.prop(node, attr_name)
def visit_Str(self, node):
self.prefix(node)
content = ast_utils.prop(node, 'content')
self.code += content if content is not None else repr(node.s)
self.suffix(node)
def check_is_elif(self, node):
try:
return ast_utils.prop(node, 'is_elif')
except AttributeError:
return False
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(ast_utils.prop(t.body[0], 'block_suffix_body'))
def test_module_suffix(self):
src = 'foo\n#bar\n\n#baz\n'
t = pasta.parse(src)
self.assertEquals(src[src.index('#bar'):], ast_utils.prop(t, 'suffix'))
