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_annotation = pytree_utils.GetNodeAnnotation(
node, pytree_utils.Annotation.SPLIT_PENALTY, default=0)
if penalty_annotation < penalty:
_SetSplitPenalty(node, penalty)
else:
for child in node.children:
RecExpression(child, first_child_leaf)
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 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, _FOO), 20)
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 _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 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 .
first = pytree_utils.FirstLeafNode(node.children[1])
if first.value != '.':
_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':
break
# 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
# 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:
stypes = pytree_utils.GetNodeAnnotation(
trailer.children[0], pytree_utils.Annotation.SUBTYPE)
if stypes and subtypes.SUBSCRIPT_BRACKET in stypes:
_SetStronglyConnected(
pytree_utils.FirstLeafNode(trailer.children[1]))
last_child_node = pytree_utils.LastLeafNode(trailer)
if last_child_node.value.strip().startswith('#'):
last_child_node = last_child_node.prev_sibling
if not (style.Get('INDENT_CLOSING_BRACKETS')
or style.Get('DEDENT_CLOSING_BRACKETS')):
last = pytree_utils.LastLeafNode(
last_child_node.prev_sibling)
if last.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])
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],
VERY_STRONGLY_CONNECTED if before else DOTTED_NAME)
_SetSplitPenalty(
node.children[1],
DOTTED_NAME if before else VERY_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 (node.children[0].type == grammar_token.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.
_RecAnnotate(node.children[-1],
pytree_utils.Annotation.SPLIT_PENALTY,
VERY_STRONGLY_CONNECTED)
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.
stypes = pytree_utils.GetNodeAnnotation(
pytree_utils.FirstLeafNode(node),
pytree_utils.Annotation.SUBTYPE)
if stypes and subtypes.SUBSCRIPT_BRACKET in stypes:
_IncreasePenalty(node, SUBSCRIPT)
# Bump up the split penalty for the first part of a subscript. We
# would rather not split there.
_IncreasePenalty(node.children[1], CONNECTED)
else:
_SetStronglyConnected(node.children[1], node.children[2])
if name == 'arglist':
_SetStronglyConnected(node.children[-1])
self.DefaultNodeVisit(node)
def testSetOnNode(self):
pytree_utils.SetNodeAnnotation(self._node, _FOO, 20)
self.assertEqual(pytree_utils.GetNodeAnnotation(self._node, _FOO), 20)
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 testGetWhenNone(self):
self.assertIsNone(pytree_utils.GetNodeAnnotation(self._leaf, _FOO))
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)