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