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)
