def testMultiple(self): pytree_utils.SetNodeAnnotation(self._leaf, _FOO, 20) pytree_utils.SetNodeAnnotation(self._leaf, _FOO1, 1) pytree_utils.SetNodeAnnotation(self._leaf, _FOO2, 2) pytree_utils.SetNodeAnnotation(self._leaf, _FOO3, 3) pytree_utils.SetNodeAnnotation(self._leaf, _FOO4, 4) pytree_utils.SetNodeAnnotation(self._leaf, _FOO5, 5) self.assertEqual(pytree_utils.GetNodeAnnotation(self._leaf, _FOO1), 1) self.assertEqual(pytree_utils.GetNodeAnnotation(self._leaf, _FOO2), 2) self.assertEqual(pytree_utils.GetNodeAnnotation(self._leaf, _FOO3), 3) self.assertEqual(pytree_utils.GetNodeAnnotation(self._leaf, _FOO4), 4) self.assertEqual(pytree_utils.GetNodeAnnotation(self._leaf, _FOO5), 5)
def testSubtype(self): pytree_utils.AppendNodeAnnotation(self._leaf, pytree_utils.Annotation.SUBTYPE, _FOO) self.assertSetEqual( pytree_utils.GetNodeAnnotation(self._leaf, pytree_utils.Annotation.SUBTYPE), {_FOO}) pytree_utils.RemoveSubtypeAnnotation(self._leaf, _FOO) self.assertSetEqual( pytree_utils.GetNodeAnnotation(self._leaf, pytree_utils.Annotation.SUBTYPE), set())
def newlines(self): """The number of newlines needed before this token.""" if style.Get('SAVE_INITIAL_BLANKLINES'): value = pytree_utils.Annotation.ORIGINAL_NEWLINES else: value = pytree_utils.Annotation.NEWLINES return pytree_utils.GetNodeAnnotation(self.node, value)
def _FindAncestorAtIndent(node, indent): """Find an ancestor of node with the given indentation. Arguments: node: node to start from. This must not be the tree root. indent: indentation string for the ancestor we're looking for. See _AnnotateIndents for more details. Returns: An ancestor node with suitable indentation. If no suitable ancestor is found, the closest ancestor to the tree root is returned. """ if node.parent.parent is None: # Our parent is the tree root, so there's nowhere else to go. return node # If the parent has an indent annotation, and it's shorter than node's # indent, this is a suitable ancestor. # The reason for "shorter" rather than "equal" is that comments may be # improperly indented (i.e. by three spaces, where surrounding statements # have either zero or two or four), and we don't want to propagate them all # the way to the root. parent_indent = pytree_utils.GetNodeAnnotation( node.parent, pytree_utils.Annotation.CHILD_INDENT) if parent_indent is not None and indent.startswith(parent_indent): return node else: # Keep looking up the tree. return _FindAncestorAtIndent(node.parent, indent)
def _AnnotateIndents(tree): """Annotate the tree with child_indent annotations. A child_indent annotation on a node specifies the indentation (as a string, like " ") of its children. It is inferred from the INDENT child of a node. Arguments: tree: root of a pytree. The pytree is modified to add annotations to nodes. Raises: RuntimeError: if the tree is malformed. """ # Annotate the root of the tree with zero indent. if tree.parent is None: pytree_utils.SetNodeAnnotation(tree, pytree_utils.Annotation.CHILD_INDENT, '') for child in tree.children: if child.type == token.INDENT: child_indent = pytree_utils.GetNodeAnnotation( tree, pytree_utils.Annotation.CHILD_INDENT) if child_indent is not None and child_indent != child.value: raise RuntimeError('inconsistent indentation for child', (tree, child)) pytree_utils.SetNodeAnnotation(tree, pytree_utils.Annotation.CHILD_INDENT, child.value) _AnnotateIndents(child)
def Visit_trailer(self, node): # pylint: disable=invalid-name # trailer ::= '(' [arglist] ')' | '[' subscriptlist ']' | '.' NAME if node.children[0].value == '.': before = style.Get('SPLIT_BEFORE_DOT') _SetSplitPenalty(node.children[0], STRONGLY_CONNECTED if before else DOTTED_NAME) _SetSplitPenalty(node.children[1], DOTTED_NAME if before else STRONGLY_CONNECTED) elif len(node.children) == 2: # Don't split an empty argument list if at all possible. _SetSplitPenalty(node.children[1], VERY_STRONGLY_CONNECTED) elif len(node.children) == 3: name = pytree_utils.NodeName(node.children[1]) if name in {'argument', 'comparison'}: # Don't split an argument list with one element if at all possible. _SetStronglyConnected(node.children[1]) if (len(node.children[1].children) > 1 and pytree_utils.NodeName(node.children[1].children[1]) == 'comp_for'): # Don't penalize splitting before a comp_for expression. _SetSplitPenalty(pytree_utils.FirstLeafNode(node.children[1]), 0) else: _SetSplitPenalty( pytree_utils.FirstLeafNode(node.children[1]), ONE_ELEMENT_ARGUMENT) elif (pytree_utils.NodeName(node.children[0]) == 'LSQB' and len(node.children[1].children) > 2 and (name.endswith('_test') or name.endswith('_expr'))): _SetStronglyConnected(node.children[1].children[0]) _SetStronglyConnected(node.children[1].children[2]) # Still allow splitting around the operator. split_before = ((name.endswith('_test') and style.Get('SPLIT_BEFORE_LOGICAL_OPERATOR')) or (name.endswith('_expr') and style.Get('SPLIT_BEFORE_BITWISE_OPERATOR'))) if split_before: _SetSplitPenalty( pytree_utils.LastLeafNode(node.children[1].children[1]), 0) else: _SetSplitPenalty( pytree_utils.FirstLeafNode(node.children[1].children[2]), 0) # Don't split the ending bracket of a subscript list. _SetVeryStronglyConnected(node.children[-1]) elif name not in { 'arglist', 'argument', 'term', 'or_test', 'and_test', 'comparison', 'atom', 'power' }: # Don't split an argument list with one element if at all possible. subtypes = pytree_utils.GetNodeAnnotation( pytree_utils.FirstLeafNode(node), pytree_utils.Annotation.SUBTYPE) if subtypes and format_token.Subtype.SUBSCRIPT_BRACKET in subtypes: _IncreasePenalty(node, SUBSCRIPT) else: _SetStronglyConnected(node.children[1], node.children[2]) if name == 'arglist': _SetStronglyConnected(node.children[-1]) self.DefaultNodeVisit(node)
def _SetTokenSubtype(self, node, subtype, force=False): """Set the token's subtype only if it's not already set.""" if force or not pytree_utils.GetNodeAnnotation( node, pytree_utils.Annotation.SUBTYPE): pytree_utils.SetNodeAnnotation(node, pytree_utils.Annotation.SUBTYPE, subtype)
def node_split_penalty(self): """Split penalty attached to the pytree node of this token. Returns: The penalty, or None if no annotation is attached. """ return pytree_utils.GetNodeAnnotation( self._node, pytree_utils.Annotation.SPLIT_PENALTY, default=0)
def Visit_comp_for(self, node): # pylint: disable=invalid-name # comp_for ::= 'for' exprlist 'in' testlist_safe [comp_iter] _AppendSubtypeRec(node, format_token.Subtype.COMP_FOR) # Mark the previous node as COMP_EXPR unless this is a nested comprehension # as these will have the outer comprehension as their previous node. attr = pytree_utils.GetNodeAnnotation(node.parent, pytree_utils.Annotation.SUBTYPE) if not attr or format_token.Subtype.COMP_FOR not in attr: _AppendSubtypeRec(node.parent.children[0], format_token.Subtype.COMP_EXPR) self.DefaultNodeVisit(node)
def FlattenRec(tree): if pytree_utils.NodeName(tree) in pytree_utils.NONSEMANTIC_TOKENS: return [] if isinstance(tree, pytree.Leaf): return [(tree.value, pytree_utils.GetNodeAnnotation( tree, pytree_utils.Annotation.SPLIT_PENALTY))] nodes = [] for node in tree.children: nodes += FlattenRec(node) return nodes
def HasDefaultOrNamedAssignSubtype(node): if isinstance(node, pytree.Leaf): if (format_token.Subtype.DEFAULT_OR_NAMED_ASSIGN in pytree_utils.GetNodeAnnotation( node, pytree_utils.Annotation.SUBTYPE, set())): return True return False has_subtype = False for child in node.children: has_subtype |= HasDefaultOrNamedAssignSubtype(child) return has_subtype
def HasSubtype(node): """Return True if the arg list has a named assign subtype.""" if isinstance(node, pytree.Leaf): if node_subtype in pytree_utils.GetNodeAnnotation( node, pytree_utils.Annotation.SUBTYPE, set()): return True return False has_subtype = False for child in node.children: if pytree_utils.NodeName(child) != 'arglist': has_subtype |= HasSubtype(child) return has_subtype
def _MustBreakBefore(prev_token, cur_token): """Return True if a line break is required before the current token.""" if prev_token.is_comment: # Must break if the previous token was a comment. return True if cur_token.is_string and prev_token.is_string: # We want consecutive strings to be on separate lines. This is a # reasonable assumption, because otherwise they should have written them # all on the same line, or with a '+'. return True return pytree_utils.GetNodeAnnotation(cur_token.node, pytree_utils.Annotation.MUST_SPLIT, default=False)
def __init__(self, node): """Constructor. Arguments: node: (pytree.Leaf) The node that's being wrapped. """ self.node = node self.next_token = None self.previous_token = None self.matching_bracket = None self.parameters = [] self.container_opening = None self.container_elements = [] self.whitespace_prefix = '' self.total_length = 0 self.split_penalty = 0 self.can_break_before = False self.must_break_before = pytree_utils.GetNodeAnnotation( node, pytree_utils.Annotation.MUST_SPLIT, default=False) self.newlines = pytree_utils.GetNodeAnnotation( node, pytree_utils.Annotation.NEWLINES) self.type = node.type self.column = node.column self.lineno = node.lineno self.name = pytree_utils.NodeName(node) self.spaces_required_before = 0 if self.is_comment: self.spaces_required_before = style.Get('SPACES_BEFORE_COMMENT') self.value = node.value if self.is_continuation: self.value = node.value.rstrip() stypes = pytree_utils.GetNodeAnnotation( node, pytree_utils.Annotation.SUBTYPE) self.subtypes = {subtypes.NONE} if not stypes else stypes self.is_pseudo = hasattr(node, 'is_pseudo') and node.is_pseudo
def HasSubtype(node): """Return True if the arg list has a named assign subtype.""" if isinstance(node, pytree.Leaf): return node_subtype in pytree_utils.GetNodeAnnotation( node, pytree_utils.Annotation.SUBTYPE, set()) for child in node.children: node_name = pytree_utils.NodeName(child) if node_name not in {'atom', 'arglist', 'power'}: if HasSubtype(child): return True return False
def RecExpression(node, first_child_leaf): if node is first_child_leaf: return if isinstance(node, pytree.Leaf): if node.value in {'(', 'for', 'if'}: return penalty = pytree_utils.GetNodeAnnotation( node, pytree_utils.Annotation.SPLIT_PENALTY, default=0) _SetSplitPenalty(node, penalty + amt) else: for child in node.children: RecExpression(child, first_child_leaf)
def Visit_funcdef(self, node): # pylint: disable=invalid-name index = self._SetBlankLinesBetweenCommentAndClassFunc(node) if _AsyncFunction(node): # Move the number of blank lines to the async keyword. num_newlines = pytree_utils.GetNodeAnnotation( node.children[0], pytree_utils.Annotation.NEWLINES) self._SetNumNewlines(node.prev_sibling, num_newlines) self._SetNumNewlines(node.children[0], None) self.last_was_decorator = False self.function_level += 1 for child in node.children[index:]: self.Visit(child) self.function_level -= 1 self.last_was_class_or_function = True
def RecArithmeticExpression(node, first_child_leaf): if node is first_child_leaf: return if isinstance(node, pytree.Leaf): if node.value in {'(', 'for', 'if'}: return penalty_annotation = pytree_utils.GetNodeAnnotation( node, pytree_utils.Annotation.SPLIT_PENALTY, default=0) if penalty_annotation < penalty: pytree_utils.SetNodeAnnotation( node, pytree_utils.Annotation.SPLIT_PENALTY, penalty) else: for child in node.children: RecArithmeticExpression(child, first_child_leaf)
def Visit_arith_expr(self, node): # pylint: disable=invalid-name # arith_expr ::= term (('+'|'-') term)* self.DefaultNodeVisit(node) _IncreasePenalty(node, ARITH_EXPR) index = 1 while index < len(node.children) - 1: child = node.children[index] if isinstance(child, pytree.Leaf) and child.value in '+-': next_node = _FirstChildNode(node.children[index + 1]) _SetSplitPenalty( next_node, pytree_utils.GetNodeAnnotation( next_node, pytree_utils.Annotation.SPLIT_PENALTY, default=0) - 100) index += 1
def _RecAnnotate(tree, annotate_name, annotate_value): """Recursively set the given annotation on all leafs of the subtree. Takes care to only increase the penalty. If the node already has a higher or equal penalty associated with it, this is a no-op. Args: tree: subtree to annotate annotate_name: name of the annotation to set annotate_value: value of the annotation to set """ for child in tree.children: _RecAnnotate(child, annotate_name, annotate_value) if isinstance(tree, pytree.Leaf): cur_annotate = pytree_utils.GetNodeAnnotation( tree, annotate_name, default=0) if cur_annotate < annotate_value: pytree_utils.SetNodeAnnotation(tree, annotate_name, annotate_value)
def _MustBreakBefore(prev_token, cur_token): """Return True if a line break is required before the current token.""" if prev_token.is_comment: # Must break if the previous token was a comment. return True if (_IsSurroundedByBrackets(cur_token) and cur_token.is_string and prev_token.is_string): # We want consecutive strings to be on separate lines. This is a # reasonable assumption, because otherwise they should have written them # all on the same line, or with a '+'. return True if style.Get('DEDENT_CLOSING_BRACKETS') and cur_token.ClosesScope(): opening = cur_token.matching_bracket trailer_length = cur_token.total_length - opening.total_length if (trailer_length > style.Get('COLUMN_LIMIT') or cur_token.lineno != opening.lineno): # Since we're already dedenting the closing bracket, let's put a newline # after the opening one so that we have more horizontal space for the # trailer. opening.next_token.must_break_before = True return True return pytree_utils.GetNodeAnnotation(cur_token.node, pytree_utils.Annotation.MUST_SPLIT, default=False)
def testSetAgain(self): pytree_utils.SetNodeAnnotation(self._leaf, _FOO, 20) self.assertEqual(pytree_utils.GetNodeAnnotation(self._leaf, _FOO), 20) pytree_utils.SetNodeAnnotation(self._leaf, _FOO, 30) self.assertEqual(pytree_utils.GetNodeAnnotation(self._leaf, _FOO), 30)
def testSetWhenNone(self): pytree_utils.SetNodeAnnotation(self._leaf, _FOO, 20) self.assertEqual(pytree_utils.GetNodeAnnotation(self._leaf, _FOO), 20)
def newlines(self): """The number of newlines needed before this token.""" return pytree_utils.GetNodeAnnotation(self._node, pytree_utils.Annotation.NEWLINES)
def subtypes(self): """Extra type information for directing formatting.""" value = pytree_utils.GetNodeAnnotation(self._node, pytree_utils.Annotation.SUBTYPE) return [Subtype.NONE] if value is None else value
def node_split_penalty(self): """Split penalty attached to the pytree node of this token.""" return pytree_utils.GetNodeAnnotation( self.node, pytree_utils.Annotation.SPLIT_PENALTY, default=0)
def must_split(self): """Return true if the token requires a split before it.""" return pytree_utils.GetNodeAnnotation( self.node, pytree_utils.Annotation.MUST_SPLIT)
def Visit_power(self, node): # pylint: disable=invalid-name,missing-docstring # power ::= atom trailer* ['**' factor] self.DefaultNodeVisit(node) # When atom is followed by a trailer, we can not break between them. # E.g. arr[idx] - no break allowed between 'arr' and '['. if (len(node.children) > 1 and pytree_utils.NodeName(node.children[1]) == 'trailer'): # children[1] itself is a whole trailer: we don't want to # mark all of it as unbreakable, only its first token: (, [ or . _SetUnbreakable(node.children[1].children[0]) # A special case when there are more trailers in the sequence. Given: # atom tr1 tr2 # The last token of tr1 and the first token of tr2 comprise an unbreakable # region. For example: foo.bar.baz(1) # We can't put breaks between either of the '.', '(', or '[' and the names # *preceding* them. prev_trailer_idx = 1 while prev_trailer_idx < len(node.children) - 1: cur_trailer_idx = prev_trailer_idx + 1 cur_trailer = node.children[cur_trailer_idx] if pytree_utils.NodeName(cur_trailer) == 'trailer': # Now we know we have two trailers one after the other prev_trailer = node.children[prev_trailer_idx] if prev_trailer.children[-1].value != ')': # Set the previous node unbreakable if it's not a function call: # atom tr1() tr2 # It may be necessary (though undesirable) to split up a previous # function call's parentheses to the next line. _SetStronglyConnected(prev_trailer.children[-1]) _SetStronglyConnected(cur_trailer.children[0]) prev_trailer_idx = cur_trailer_idx else: break # We don't want to split before the last ')' of a function call. This also # takes care of the special case of: # atom tr1 tr2 ... trn # where the 'tr#' are trailers that may end in a ')'. for trailer in node.children[1:]: if pytree_utils.NodeName(trailer) != 'trailer': break if trailer.children[0].value in '([': if len(trailer.children) > 2: subtypes = pytree_utils.GetNodeAnnotation( trailer.children[0], pytree_utils.Annotation.SUBTYPE) if subtypes and format_token.Subtype.SUBSCRIPT_BRACKET in subtypes: _SetStronglyConnected( _FirstChildNode(trailer.children[1])) last_child_node = _LastChildNode(trailer) if last_child_node.value.strip().startswith('#'): last_child_node = last_child_node.prev_sibling if not style.Get('DEDENT_CLOSING_BRACKETS'): if _LastChildNode( last_child_node.prev_sibling).value != ',': if last_child_node.value == ']': _SetUnbreakable(last_child_node) else: _SetSplitPenalty(last_child_node, VERY_STRONGLY_CONNECTED) else: # If the trailer's children are '()', then make it a strongly # connected region. It's sometimes necessary, though undesirable, to # split the two. _SetStronglyConnected(trailer.children[-1]) # If the original source has a "builder" style calls, then we should allow # the reformatter to retain that. _AllowBuilderStyleCalls(node)
def _DecrementSplitPenalty(node, amt): penalty = pytree_utils.GetNodeAnnotation( node, pytree_utils.Annotation.SPLIT_PENALTY, default=amt) penalty = penalty - amt if amt < penalty else 0 _SetSplitPenalty(node, penalty)
def testSetOnNode(self): pytree_utils.SetNodeAnnotation(self._node, _FOO, 20) self.assertEqual(pytree_utils.GetNodeAnnotation(self._node, _FOO), 20)