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
val = fmt.get(node, attr_name)
if (val is None or deps and any(
getattr(node, dep, None) != fmt.get(node, dep + '__src')
for dep in deps)):
val = default
self.code += val if val is not None else ''
def test(self):
with open(input_file, 'r') as handle:
src = handle.read()
t = ast_utils.parse(src, py_ver)
annotator = annotate.get_ast_annotator(py_ver)(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 pasta.ast_walk(t, py_ver)) + '\n'
# If specified, write the golden data instead of checking it
if getattr(self, 'generate_goldens', False):
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 replace_child(parent, node, replace_with):
"""Replace a node's child with another node while preserving formatting.
Arguments:
parent: (ast.AST) Parent node to replace a child of.
node: (ast.AST) Child node to replace.
replace_with: (ast.AST) New child node.
"""
# TODO(soupytwist): Don't refer to the formatting dict directly
if hasattr(node, fmt.PASTA_DICT):
fmt.set(replace_with, 'prefix', fmt.get(node, 'prefix'))
fmt.set(replace_with, 'suffix', fmt.get(node, 'suffix'))
for field in parent._fields:
field_val = getattr(parent, field, None)
if field_val == node:
setattr(parent, field, replace_with)
return
elif isinstance(field_val, list):
try:
field_val[field_val.index(node)] = replace_with
return
except ValueError:
pass
raise errors.InvalidAstError('Node %r is not a child of %r' %
(node, parent))
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(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 open(golden_file, 'r') as f:
golden = f.read()
except IOError:
self.fail('Missing golden data.')
self.assertMultiLineEqual(golden, result)
def test_indent_multiline_string_with_newline(self):
src = textwrap.dedent('''\
class A:
"""Doc\n
string."""
pass
''')
t = pasta.parse(src, py_ver)
docstring, pass_stmt = t.body[0].body
self.assertEqual(' ', fmt.get(docstring, 'indent'))
self.assertEqual(' ', fmt.get(pass_stmt, 'indent'))
def close_scope(self,
node,
prefix_attr='prefix',
suffix_attr='suffix',
trailing_comma=False,
single_paren=False):
"""Close a parenthesized scope on the given node, if one is open."""
# Ensures the prefix + suffix are not None
if fmt.get(node, prefix_attr) is None:
fmt.set(node, prefix_attr, '')
if fmt.get(node, suffix_attr) is None:
fmt.set(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
encountered_paren = False
result = ''
for tok in self.takewhile(
lambda t: t.type in FORMATTING_TOKENS or t.src in symbols):
# Consume all space up to this token
result += self._space_between(prev_loc, tok)
if tok.src == ')' and single_paren and encountered_paren:
self.rewind()
parsed_to_i = self._i
parsed_to_loc = tok.start
fmt.append(node, suffix_attr, result)
break
# Consume the token itself
result += tok.src
if tok.src == ')':
# Close out the open scope
encountered_paren = True
self._scope_stack.pop()
fmt.prepend(node, prefix_attr, self._parens.pop())
fmt.append(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 get_last_child(node):
"""Get the last child node of a block statement.
The input must be a block statement (e.g. ast.For, ast.With, etc).
Examples:
1. with first():
second()
last()
2. try:
first()
except:
second()
finally:
last()
In both cases, the last child is the node for `last`.
"""
if isinstance(node, ast27.Module) or isinstance(node, ast3.Module):
try:
return node.body[-1]
except IndexError:
return None
if isinstance(node, ast27.If) or isinstance(node, ast3.If):
if (len(node.orelse) == 1 and isinstance(node.orelse[0],
(ast27.If, ast3.If)) and
fmt.get(node.orelse[0], 'is_elif')):
return get_last_child(node.orelse[0])
if node.orelse:
return node.orelse[-1]
elif isinstance(node, ast27.With) or isinstance(node, ast3.With):
if (len(node.body) == 1 and isinstance(node.body[0],
(ast27.With, ast3.With)) and
fmt.get(node.body[0], 'is_continued')):
return get_last_child(node.body[0])
elif isinstance(node, ast3.Try):
if node.finalbody:
return node.finalbody[-1]
if node.orelse:
return node.orelse[-1]
elif isinstance(node, ast27.TryFinally):
if node.finalbody:
return node.finalbody[-1]
elif isinstance(node, ast27.TryExcept):
if node.orelse:
return node.orelse[-1]
if node.handlers:
return get_last_child(node.handlers[-1])
return node.body[-1]
def visit_Module(self, node):
self.prefix(node)
bom = fmt.get(node, 'bom')
if bom is not None:
self.code += bom
self.generic_visit(node)
self.suffix(node)
def optional_token(self, node, attr_name, token_val,
allow_whitespace_prefix=False, default=False):
del allow_whitespace_prefix
value = fmt.get(node, attr_name)
if value is None and default:
value = token_val
self.code += value or ''
def test_fstring(self):
src = 'f"a {b} c d {e}"'
t = pasta.parse(src, py_ver)
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_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 visit_Constant(self, node):
self.prefix(node)
if node.value is Ellipsis:
content = '...'
else:
content = fmt.get(node, 'content')
self.code += content if content is not None else repr(node.s)
self.suffix(node)
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_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_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, py_ver)
if_node = t.body[0]
b = if_node.body[0]
self.assertEqual(' ', fmt.get(b, 'indent_diff'))
def wrapped(self, node, *args, **kwargs):
self.prefix(node, default=self._indent)
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 = (fmt.get(last_child, 'prefix') or '\n').splitlines()[-1]
self.block_suffix(node, indent)
else:
self.suffix(node, comment=True)
def test_tab_below_spaces(self):
for num_spaces in range(1, 8):
t = pasta.parse(
textwrap.dedent("""\
if a:
{WS}if b:
{ONETAB}c
""").format(ONETAB='\t', WS=' ' * num_spaces), py_ver)
node_c = t.body[0].body[0].body[0]
self.assertEqual(fmt.get(node_c, 'indent_diff'),
' ' * (8 - num_spaces))
def test_indent_extra_newlines_with_comment(self):
src = textwrap.dedent('''\
if a:
#not here
b
''')
t = pasta.parse(src)
if_node = t.body[0]
b = if_node.body[0]
self.assertEqual(' ', fmt.get(b, 'indent_diff'))
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]) == ast27.Pass
or type(val[0]) == ast3.Pass)
for children_attr, open_block in test_cases:
src = src_tpl.format(open_block=open_block)
t = pasta.parse(src, py_ver)
node_finder = ast_utils.get_find_node_visitor(
lambda node: is_node_for_suffix(node, children_attr),
py_ver)
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, py_ver))
def indented(self, node, children_attr):
prev_indent = self._indent
prev_indent_diff = self._indent_diff
new_diff = fmt.get(node, 'indent')
if new_diff is None:
new_diff = ' '
self._indent_diff = new_diff
self._indent = prev_indent + self._indent_diff
for child in getattr(node, children_attr):
yield child
self.attr(node, 'block_suffix_%s' % children_attr, [])
self._indent = prev_indent
self._indent_diff = prev_indent_diff
def _modify_function_name(func_def_node, new_func_name):
"""Modify function name"""
if not isinstance(func_def_node, ast.FunctionDef):
raise NodeTypeNotSupport('It is not ast.FunctionDef node type.')
old_func_name = func_def_node.name
func_def_node.name = new_func_name
# Modify formatting information stored by pasta
old_function_def = fmt.get(func_def_node, 'function_def')
if old_function_def:
new_function_def = old_function_def.replace(
old_func_name, new_func_name)
fmt.set(func_def_node, 'function_def', new_function_def)
fmt.set(func_def_node, 'name__src', new_func_name)
def visit_JoinedStr(self, node):
self.prefix(node)
content = fmt.get(node, 'content')
if content is None:
parts = []
for val in node.values:
if isinstance(val, ast.Str):
parts.append(val.s)
else:
parts.append(fstring_utils.placeholder(len(parts)))
content = repr(''.join(parts))
values = [to_str(v) for v in fstring_utils.get_formatted_values(node)]
self.code += fstring_utils.perform_replacements(content, values)
self.suffix(node)
def to_str(tree):
"""Convenient function to get the python source for an AST."""
p = Printer()
# Detect the most prevalent indentation style in the file and use it when
# printing indented nodes which don't have formatting data.
seen_indent_diffs = collections.defaultdict(lambda: 0)
for node in ast.walk(tree):
seen_indent_diffs[fmt.get(node, 'indent_diff', '')] += 1
del seen_indent_diffs['']
if seen_indent_diffs:
indent_diff, _ = max(six.iteritems(seen_indent_diffs),
key=lambda tup: tup[1] if tup[0] else -1)
p.set_default_indent_diff(indent_diff)
p.visit(tree)
return p.code
def visit_If(self, node):
tok = 'elif' if fmt.get(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 self.indented(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])):
fmt.set(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 self.indented(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('', 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.assertEquals(src[src.index('#bar'):], fmt.get(t, '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', fmt.get(t.body[1], 'prefix'))
self.assertEqual('', fmt.get(t.body[1], 'suffix'))
