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 _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 _AdjustSplitPenalty(line):
"""Visit the node and adjust the split penalties if needed.
A token shouldn't be split if it's not within a bracket pair. Mark any token
that's not within a bracket pair as "unbreakable".
Arguments:
line: (LogicalLine) An logical line.
"""
bracket_level = 0
for index, token in enumerate(line.tokens):
if index and not bracket_level:
pytree_utils.SetNodeAnnotation(
token.node, pytree_utils.Annotation.SPLIT_PENALTY,
split_penalty.UNBREAKABLE)
if token.value in _OPENING_BRACKETS:
bracket_level += 1
elif token.value in _CLOSING_BRACKETS:
bracket_level -= 1
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 _SetSplitPenalty(node, penalty):
pytree_utils.SetNodeAnnotation(node, pytree_utils.Annotation.SPLIT_PENALTY,
penalty)
def _SetNumNewlines(node, num_newlines):
pytree_utils.SetNodeAnnotation(node, pytree_utils.Annotation.NEWLINES,
num_newlines)
def _SetMustSplitOnFirstLeaf(node):
"""Set the "must split" annotation on the first leaf node."""
pytree_utils.SetNodeAnnotation(pytree_utils.FirstLeafNode(node),
pytree_utils.Annotation.MUST_SPLIT, True)
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)
