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