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(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))