def visit(self, node):
try:
fmt.set(node, 'indent', self._indent)
fmt.set(node, 'indent_diff', self._indent_diff)
super(AstAnnotator, self).visit(node)
except (TypeError, ValueError, IndexError, KeyError) as e:
raise AnnotationError(e)
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 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 _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 attr(self, node, attr_name, attr_vals, deps=None, default=None):
"""Parses some source and sets an attribute on the given node.
Stores some arbitrary formatting information on the node. This takes a list
attr_vals which tell what parts of the source to parse. The result of each
function is concatenated onto the formatting data, and strings in this list
are a shorthand to look for an exactly matching token.
For example:
self.attr(node, 'foo', ['(', self.ws, 'Hello, world!', self.ws, ')'],
deps=('s',), default=node.s)
is a rudimentary way to parse a parenthesized string. After running this,
the matching source code for this node will be stored in its formatting
dict under the key 'foo'. The result might be `(\n 'Hello, world!'\n)`.
This also keeps track of the current value of each of the dependencies.
In the above example, we would have looked for the string 'Hello, world!'
because that's the value of node.s, however, when we print this back, we
want to know if the value of node.s has changed since this time. If any of
the dependent values has changed, the default would be used instead.
Arguments:
node: (ast.AST) An AST node to attach formatting information to.
attr_name: (string) Name to store the formatting information under.
attr_vals: (list of functions/strings) Each item is either a function
that parses some source and return a string OR a string to match
exactly (as a token).
deps: (optional, set of strings) Attributes of the node which attr_vals
depends on.
default: (string) Unused here.
"""
del default # unused
if deps:
for dep in deps:
fmt.set(node, dep + '__src', getattr(node, dep, None))
attr_parts = []
for attr_val in attr_vals:
if isinstance(attr_val, six.string_types):
attr_parts.append(self.token(attr_val))
else:
attr_parts.append(attr_val())
fmt.set(node, attr_name, ''.join(attr_parts))
def indented(self, node, children_attr):
"""Generator white annotates child nodes with their indentation level."""
children = getattr(node, children_attr)
cur_loc = self.tokens._loc
next_loc = self.tokens.peek_non_whitespace().start
# Special case: if the children are on the same line, then there is no
# indentation level to track.
if cur_loc[0] == next_loc[0]:
indent_diff = self._indent_diff
self._indent_diff = None
for child in children:
yield child
self._indent_diff = indent_diff
return
prev_indent = self._indent
prev_indent_diff = self._indent_diff
# Find the indent level of the first child
indent_token = self.tokens.peek_conditional(
lambda t: t.type == token_generator.TOKENS.INDENT)
new_indent = indent_token.src
if (not new_indent.startswith(prev_indent) or
len(new_indent) <= len(prev_indent)):
raise AnnotationError(
'Indent detection failed (line %d); inner indentation level is not '
'more than the outer indentation.' % lineno)
# Set the indent level to the child's indent and iterate over the children
self._indent = new_indent
self._indent_diff = new_indent[len(prev_indent):]
for child in children:
yield child
# Store the suffix at this indentation level, which could be many lines
fmt.set(node, 'block_suffix_%s' % children_attr,
self.tokens.block_whitespace(self._indent))
# Dedent back to the previous level
self._indent = prev_indent
self._indent_diff = prev_indent_diff
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 block_suffix(self, node, indent_level): fmt.set(node, 'suffix', self.tokens.block_whitespace(indent_level))