def Visit_atom(self, node): # pylint: disable=invalid-name
# atom ::= ('(' [yield_expr|testlist_gexp] ')'
# '[' [listmaker] ']' |
# '{' [dictsetmaker] '}')
self.DefaultNodeVisit(node)
if node.children[0].value == '(':
if node.children[0].lineno == node.children[-1].lineno:
self._SetExpressionPenalty(node, CONTIGUOUS_LIST)
if node.children[-1].value == ')':
if pytree_utils.NodeName(node.parent) == 'if_stmt':
pytree_utils.SetNodeAnnotation(
node.children[-1],
pytree_utils.Annotation.SPLIT_PENALTY, UNBREAKABLE)
else:
pytree_utils.SetNodeAnnotation(
node.children[-1],
pytree_utils.Annotation.SPLIT_PENALTY,
STRONGLY_CONNECTED)
elif node.children[0].value in '[{':
# Keep empty containers together if we can.
lbracket = node.children[0]
rbracket = node.children[-1]
if len(node.children) == 2:
self._SetUnbreakable(node.children[-1])
elif (rbracket.value in ']}'
and lbracket.get_lineno() == rbracket.get_lineno()
and rbracket.column - lbracket.column <
style.Get('COLUMN_LIMIT')):
self._SetExpressionPenalty(node, CONTIGUOUS_LIST)
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_comparison(self, node): # pylint: disable=invalid-name
# comparison ::= expr (comp_op expr)*
self.DefaultNodeVisit(node)
if len(node.children) == 3 and _StronglyConnectedCompOp(node):
pytree_utils.SetNodeAnnotation(
_FirstChildNode(node.children[1]),
pytree_utils.Annotation.SPLIT_PENALTY, STRONGLY_CONNECTED)
pytree_utils.SetNodeAnnotation(
_FirstChildNode(node.children[2]),
pytree_utils.Annotation.SPLIT_PENALTY, STRONGLY_CONNECTED)
else:
self._SetExpressionPenalty(node, COMPARISON_EXPRESSION)
def Visit_trailer(self, node): # pylint: disable=invalid-name
# trailer ::= '(' [arglist] ')' | '[' subscriptlist ']' | '.' NAME
self.DefaultNodeVisit(node)
if node.children[0].value == '.':
self._SetUnbreakableOnChildren(node)
pytree_utils.SetNodeAnnotation(
node.children[1], pytree_utils.Annotation.SPLIT_PENALTY,
DOTTED_NAME)
elif len(node.children) == 2:
# Don't split an empty argument list if at all possible.
pytree_utils.SetNodeAnnotation(
node.children[1], pytree_utils.Annotation.SPLIT_PENALTY,
VERY_STRONGLY_CONNECTED)
elif len(node.children) == 3:
name = pytree_utils.NodeName(node.children[1])
if name == 'power':
if pytree_utils.NodeName(
node.children[1].children[0]) != 'atom':
# Don't split an argument list with one element if at all possible.
self._SetStronglyConnected(node.children[1],
node.children[2])
pytree_utils.SetNodeAnnotation(
_FirstChildNode(node.children[1]),
pytree_utils.Annotation.SPLIT_PENALTY,
ONE_ELEMENT_ARGUMENT)
elif (pytree_utils.NodeName(node.children[0]) == 'LSQB'
and (name.endswith('_test') or name.endswith('_expr'))):
self._SetStronglyConnected(node.children[1].children[0])
self._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:
pytree_utils.SetNodeAnnotation(
_LastChildNode(node.children[1].children[1]),
pytree_utils.Annotation.SPLIT_PENALTY, 0)
else:
pytree_utils.SetNodeAnnotation(
_FirstChildNode(node.children[1].children[2]),
pytree_utils.Annotation.SPLIT_PENALTY, 0)
# Don't split the ending bracket of a subscript list.
self._SetVeryStronglyConnected(node.children[-1])
elif name not in {
'arglist', 'argument', 'term', 'or_test', 'and_test',
'comparison', 'atom'
}:
# Don't split an argument list with one element if at all possible.
self._SetStronglyConnected(node.children[1], node.children[2])
def Visit_funcdef(self, node): # pylint: disable=invalid-name
# funcdef ::= 'def' NAME parameters ['->' test] ':' suite
#
# Can't break before the function name and before the colon. The parameters
# are handled by child iteration.
colon_idx = 1
while pytree_utils.NodeName(node.children[colon_idx]) == 'simple_stmt':
colon_idx += 1
self._SetUnbreakable(node.children[colon_idx])
arrow_idx = -1
while colon_idx < len(node.children):
if isinstance(node.children[colon_idx], pytree.Leaf):
if node.children[colon_idx].value == ':':
break
if node.children[colon_idx].value == '->':
arrow_idx = colon_idx
colon_idx += 1
self._SetUnbreakable(node.children[colon_idx])
self.DefaultNodeVisit(node)
if arrow_idx > 0:
pytree_utils.SetNodeAnnotation(
_LastChildNode(node.children[arrow_idx - 1]),
pytree_utils.Annotation.SPLIT_PENALTY, 0)
self._SetUnbreakable(node.children[arrow_idx])
self._SetStronglyConnected(node.children[arrow_idx + 1])
def Split(self):
"""Split the line at semicolons."""
if not self.has_semicolon or self.disable:
return [self]
uwlines = []
uwline = UnwrappedLine(self.depth)
for tok in self._tokens:
if tok.value == ';':
uwlines.append(uwline)
uwline = UnwrappedLine(self.depth)
else:
uwline.AppendToken(tok)
if uwline.tokens:
uwlines.append(uwline)
for uwline in uwlines:
pytree_utils.SetNodeAnnotation(uwline.first.node,
pytree_utils.Annotation.MUST_SPLIT,
True)
uwline.first.previous_token = None
uwline.last.next_token = None
return uwlines
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 Visit_comp_if(self, node): # pylint: disable=invalid-name
# comp_if ::= 'if' old_test [comp_iter]
pytree_utils.SetNodeAnnotation(node.children[0],
pytree_utils.Annotation.SPLIT_PENALTY,
0)
self._SetStronglyConnected(*node.children[1:])
self.DefaultNodeVisit(node)
def Visit_comp_for(self, node): # pylint: disable=invalid-name
# comp_for ::= 'for' exprlist 'in' testlist_safe [comp_iter]
pytree_utils.SetNodeAnnotation(_FirstChildNode(node),
pytree_utils.Annotation.SPLIT_PENALTY,
0)
self._SetStronglyConnected(*node.children[1:])
self.DefaultNodeVisit(node)
def move_lines_to_index(uwline_index_to, lineno, uwlines, lines):
"""Method moves all lines in the list to the proper index of uwlines and
update lineno on these lines. This is useful when you want to change the
order of code lines. But note: it is not updating lineno on other lines
@:returns positions (indexes) from original source where
lines are taken from
"""
# saving positions of imports here, that will be used for restoring 'lineno'
lineno_where_line_was_taken_from = list()
for line in lines:
lineno_where_line_was_taken_from.append(line.lineno)
for token in line.tokens:
# here we will restore correct lineno for moved lines
token.node.lineno = lineno
# hack to remove newlines between imports that we moved to top
pytree_utils.SetNodeAnnotation(token.node,
pytree_utils.Annotation.NEWLINES, 0)
lineno += get_lineno_delta(token)
# need to update lineno on import lines to have consistency
lineno += 1
# filtering moved values and removing them from uwlines
uwlines[:] = [line for line in uwlines if line not in lines]
uwlines[uwline_index_to:uwline_index_to] = lines
return lineno_where_line_was_taken_from
def _SetMustSplitOnFirstLeaf(node):
"""Set the "must split" annotation on the first leaf node."""
def FindFirstLeaf(node):
if isinstance(node, pytree.Leaf):
return node
return FindFirstLeaf(node.children[0])
pytree_utils.SetNodeAnnotation(FindFirstLeaf(node),
pytree_utils.Annotation.MUST_SPLIT, True)
def Visit_import_as_names(self, node): # pyline: disable=invalid-name
# import_as_names ::= import_as_name (',' import_as_name)* [',']
self.DefaultNodeVisit(node)
prev_child = None
for child in node.children:
if (prev_child and isinstance(prev_child, pytree.Leaf)
and prev_child.value == ','):
pytree_utils.SetNodeAnnotation(
child, pytree_utils.Annotation.SPLIT_PENALTY,
style.Get('SPLIT_PENALTY_IMPORT_NAMES'))
prev_child = child
def Visit_and_test(self, node): # pylint: disable=invalid-name
# and_test ::= not_test ('and' not_test)*
self.DefaultNodeVisit(node)
_SetExpressionPenalty(node, AND_TEST)
index = 1
while index + 1 < len(node.children):
if style.Get('SPLIT_BEFORE_LOGICAL_OPERATOR'):
pytree_utils.SetNodeAnnotation(
_FirstChildNode(node.children[index]),
pytree_utils.Annotation.SPLIT_PENALTY, 0)
pytree_utils.SetNodeAnnotation(
_FirstChildNode(node.children[index + 1]),
pytree_utils.Annotation.SPLIT_PENALTY, STRONGLY_CONNECTED)
else:
pytree_utils.SetNodeAnnotation(
_FirstChildNode(node.children[index]),
pytree_utils.Annotation.SPLIT_PENALTY, STRONGLY_CONNECTED)
pytree_utils.SetNodeAnnotation(
_FirstChildNode(node.children[index + 1]),
pytree_utils.Annotation.SPLIT_PENALTY, 0)
index += 2
def Visit_atom(self, node): # pylint: disable=invalid-name
# atom ::= ('(' [yield_expr|testlist_gexp] ')'
# '[' [listmaker] ']' |
# '{' [dictsetmaker] '}')
self.DefaultNodeVisit(node)
if node.children[0].value == '(':
if node.children[-1].value == ')':
if pytree_utils.NodeName(node.parent) == 'if_stmt':
pytree_utils.SetNodeAnnotation(
node.children[-1],
pytree_utils.Annotation.SPLIT_PENALTY, UNBREAKABLE)
else:
pytree_utils.SetNodeAnnotation(
node.children[-1],
pytree_utils.Annotation.SPLIT_PENALTY,
COMPARISON_EXPRESSION)
elif node.children[0].value in '[{':
# Keep empty containers together if we can.
lbracket = node.children[0]
rbracket = node.children[-1]
if len(node.children) == 2:
_SetUnbreakable(node.children[-1])
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_atom(self, node): # pylint: disable=invalid-name
# atom ::= ('(' [yield_expr|testlist_gexp] ')'
# '[' [listmaker] ']' |
# '{' [dictsetmaker] '}')
self.DefaultNodeVisit(node)
if node.children[0].value == '(':
if node.children[0].lineno == node.children[-1].lineno:
self._SetExpressionPenalty(node, CONTIGUOUS_LIST)
if (pytree_utils.NodeName(node.parent) == 'if_stmt'
and node.children[-1].value == ')'):
pytree_utils.SetNodeAnnotation(
node.children[-1], pytree_utils.Annotation.SPLIT_PENALTY,
UNBREAKABLE)
elif node.children[0].value in '[{':
# Keep empty containers together if we can.
if len(node.children) == 2:
self._SetStronglyConnected(node.children[-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 _AllowBuilderStyleCalls(node):
"""Allow splitting before '.' if it's a builder style function call."""
def RecGetLeaves(node):
if isinstance(node, pytree.Leaf):
return [node]
children = []
for child in node.children:
children += RecGetLeaves(child)
return children
list_of_children = RecGetLeaves(node)
prev_child = None
for child in list_of_children:
if child.value == '.':
if prev_child.lineno != child.lineno:
pytree_utils.SetNodeAnnotation(
child, pytree_utils.Annotation.SPLIT_PENALTY, 0)
prev_child = child
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 Visit_trailer(self, node): # pylint: disable=invalid-name
# trailer ::= '(' [arglist] ')' | '[' subscriptlist ']' | '.' NAME
self.DefaultNodeVisit(node)
if node.children[0].value == '.':
self._SetUnbreakableOnChildren(node)
pytree_utils.SetNodeAnnotation(
node.children[1], pytree_utils.Annotation.SPLIT_PENALTY,
DOTTED_NAME)
elif len(node.children) == 2:
# Don't split an empty argument list if at all possible.
self._SetStronglyConnected(node.children[1])
elif len(node.children) == 3:
if (pytree_utils.NodeName(node.children[1])
not in {'arglist', 'argument', 'term'}):
# Don't split an argument list with one element if at all possible.
self._SetStronglyConnected(node.children[1], node.children[2])
if pytree_utils.NodeName(node.children[-1]) == 'RSQB':
# Don't split the ending bracket of a subscript list.
self._SetStronglyConnected(node.children[-1])
def _AdjustSplitPenalty(uwline):
"""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:
uwline: (UnwrappedLine) An unwrapped line.
"""
bracket_level = 0
for index, token in enumerate(uwline.tokens):
if index and not bracket_level:
pytree_utils.SetNodeAnnotation(token.node,
pytree_utils.Annotation.SPLIT_PENALTY,
split_penalty.UNBREAKABLE)
if token.value in pytree_utils.OPENING_BRACKETS:
bracket_level += 1
elif token.value in pytree_utils.CLOSING_BRACKETS:
bracket_level -= 1
def _SetNumNewlines(self, node, num_newlines):
pytree_utils.SetNodeAnnotation(node, pytree_utils.Annotation.NEWLINES,
num_newlines)
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 .
self._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.
self._SetStronglyConnected(prev_trailer.children[-1])
self._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:
self._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 == ']':
self._SetUnbreakable(last_child_node)
else:
pytree_utils.SetNodeAnnotation(
last_child_node, pytree_utils.Annotation.SPLIT_PENALTY,
VERY_STRONGLY_CONNECTED)
if _FirstChildNode(trailer).lineno == last_child_node.lineno:
# If the trailer was originally on one line, then try to keep it
# like that.
self._SetExpressionPenalty(trailer, CONTIGUOUS_LIST)
else:
# If the trailer's children are '()', then make it a strongly
# connected region. It's sometimes necessary, though undesirable, to
# split the two.
self._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 _CalculateNumberOfNewlines(first_token, indent_depth, prev_uwline,
final_lines):
"""Calculate the number of newlines we need to add.
Arguments:
first_token: (format_token.FormatToken) The first token in the unwrapped
line.
indent_depth: (int) The line's indentation depth.
prev_uwline: (list of unwrapped_line.UnwrappedLine) The unwrapped line
previous to this line.
final_lines: (list of unwrapped_line.UnwrappedLine) The unwrapped lines
that have already been processed.
Returns:
The number of newlines needed before the first token.
"""
# TODO(morbo): Special handling for imports.
# TODO(morbo): Create a knob that can tune these.
if prev_uwline is None:
# The first line in the file. Don't add blank lines.
# FIXME(morbo): Is this correct?
if first_token.newlines is not None:
pytree_utils.SetNodeAnnotation(first_token.node,
pytree_utils.Annotation.NEWLINES,
None)
return 0
if first_token.is_docstring:
if (prev_uwline.first.value == 'class'
and style.Get('BLANK_LINE_BEFORE_CLASS_DOCSTRING')):
# Enforce a blank line before a class's docstring.
return ONE_BLANK_LINE
elif (prev_uwline.first.value.startswith('#')
and style.Get('BLANK_LINE_BEFORE_MODULE_DOCSTRING')):
# Enforce a blank line before a module's docstring.
return ONE_BLANK_LINE
# The docstring shouldn't have a newline before it.
return NO_BLANK_LINES
prev_last_token = prev_uwline.last
if prev_last_token.is_docstring:
if (not indent_depth
and first_token.value in {'class', 'def', 'async'}):
# Separate a class or function from the module-level docstring with
# appropriate number of blank lines.
return 1 + style.Get('BLANK_LINES_AROUND_TOP_LEVEL_DEFINITION')
if _NoBlankLinesBeforeCurrentToken(prev_last_token.value, first_token,
prev_last_token):
return NO_BLANK_LINES
else:
return ONE_BLANK_LINE
if first_token.value in {'class', 'def', 'async', '@'}:
# TODO(morbo): This can go once the blank line calculator is more
# sophisticated.
if not indent_depth:
# This is a top-level class or function.
is_inline_comment = prev_last_token.whitespace_prefix.count(
'\n') == 0
if (not prev_uwline.disable and prev_last_token.is_comment
and not is_inline_comment):
# This token follows a non-inline comment.
if _NoBlankLinesBeforeCurrentToken(prev_last_token.value,
first_token,
prev_last_token):
# Assume that the comment is "attached" to the current line.
# Therefore, we want two blank lines before the comment.
index = len(final_lines) - 1
while index > 0:
if not final_lines[index - 1].is_comment:
break
index -= 1
if final_lines[index - 1].first.value == '@':
final_lines[index].first.AdjustNewlinesBefore(
NO_BLANK_LINES)
else:
prev_last_token.AdjustNewlinesBefore(1 + style.Get(
'BLANK_LINES_AROUND_TOP_LEVEL_DEFINITION'))
if first_token.newlines is not None:
pytree_utils.SetNodeAnnotation(
first_token.node, pytree_utils.Annotation.NEWLINES,
None)
return NO_BLANK_LINES
elif _IsClassOrDef(prev_uwline):
if not style.Get('BLANK_LINE_BEFORE_NESTED_CLASS_OR_DEF'):
pytree_utils.SetNodeAnnotation(
first_token.node, pytree_utils.Annotation.NEWLINES, None)
return NO_BLANK_LINES
# Calculate how many newlines were between the original lines. We want to
# retain that formatting if it doesn't violate one of the style guide rules.
if first_token.is_comment:
first_token_lineno = first_token.lineno - first_token.value.count('\n')
else:
first_token_lineno = first_token.lineno
prev_last_token_lineno = prev_last_token.lineno
if prev_last_token.is_multiline_string:
prev_last_token_lineno += prev_last_token.value.count('\n')
if first_token_lineno - prev_last_token_lineno > 1:
return ONE_BLANK_LINE
return NO_BLANK_LINES
def testSetWhenNone(self): pytree_utils.SetNodeAnnotation(self._leaf, _FOO, 20) self.assertEqual(pytree_utils.GetNodeAnnotation(self._leaf, _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 testSetOnNode(self): pytree_utils.SetNodeAnnotation(self._node, _FOO, 20) self.assertEqual(pytree_utils.GetNodeAnnotation(self._node, _FOO), 20)
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 _CalculateNumberOfNewlines(first_token, indent_depth, prev_uwline):
"""Calculate the number of newlines we need to add.
Arguments:
first_token: (format_token.FormatToken) The first token in the unwrapped
line.
indent_depth: (int) The line's indentation depth.
prev_uwline: (list of unwrapped_line.UnwrappedLine) The unwrapped line
previous to this line.
Returns:
The number of newlines needed before the first token.
"""
# TODO(morbo): Special handling for imports.
# TODO(morbo): Create a knob that can tune these.
if prev_uwline is None:
# The first line in the file. Don't add blank lines.
# FIXME(morbo): Is this correct?
if first_token.newlines is not None:
pytree_utils.SetNodeAnnotation(first_token.GetPytreeNode(),
pytree_utils.Annotation.NEWLINES, None)
return 0
if first_token.is_docstring:
# The docstring shouldn't have a newline before it.
# TODO(morbo): Add a knob to adjust this.
return NO_BLANK_LINES
prev_last_token = prev_uwline.last
if prev_last_token.is_docstring:
if not indent_depth and first_token.value in {'class', 'def'}:
# Separate a class or function from the module-level docstring with two
# blank lines.
return TWO_BLANK_LINES
if _NoBlankLinesBeforeCurrentToken(prev_last_token.value, first_token,
prev_last_token):
return NO_BLANK_LINES
else:
return ONE_BLANK_LINE
if first_token.value in {'class', 'def', '@'}:
# TODO(morbo): This can go once the blank line calculator is more
# sophisticated.
if not indent_depth:
# This is a top-level class or function.
is_inline_comment = prev_last_token.whitespace_prefix.count('\n') == 0
if prev_last_token.is_comment and not is_inline_comment:
# This token follows a non-inline comment.
if _NoBlankLinesBeforeCurrentToken(prev_last_token.value, first_token,
prev_last_token):
# Assume that the comment is "attached" to the current line.
# Therefore, we want two blank lines before the comment.
prev_last_token.AdjustNewlinesBefore(TWO_BLANK_LINES)
if first_token.newlines is not None:
pytree_utils.SetNodeAnnotation(first_token.GetPytreeNode(),
pytree_utils.Annotation.NEWLINES,
None)
return NO_BLANK_LINES
elif prev_uwline.first.value in {'class', 'def'}:
if not style.Get('BLANK_LINE_BEFORE_NESTED_CLASS_OR_DEF'):
pytree_utils.SetNodeAnnotation(first_token.node,
pytree_utils.Annotation.NEWLINES, None)
return NO_BLANK_LINES
# Calculate how many newlines were between the original lines. We want to
# retain that formatting if it doesn't violate one of the style guide rules.
if first_token.is_comment:
first_token_lineno = first_token.lineno - first_token.value.count('\n')
else:
first_token_lineno = first_token.lineno
if first_token_lineno - prev_last_token.lineno > 1:
return ONE_BLANK_LINE
return NO_BLANK_LINES
def _SetSplitPenalty(node, penalty):
pytree_utils.SetNodeAnnotation(node, pytree_utils.Annotation.SPLIT_PENALTY,
penalty)
