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