def testPrintFunctionToTree(self):
tree = pytree_utils.ParseCodeToTree(
'print("hello world", file=sys.stderr)\n')
self.assertEqual('file_input', pytree_utils.NodeName(tree))
self.assertEqual(2, len(tree.children))
self.assertEqual('simple_stmt',
pytree_utils.NodeName(tree.children[0]))
def testParseCodeToTree(self):
# Since ParseCodeToTree is a thin wrapper around underlying lib2to3
# functionality, only a sanity test here...
tree = pytree_utils.ParseCodeToTree('foo = 2\n')
self.assertEqual('file_input', pytree_utils.NodeName(tree))
self.assertEqual(2, len(tree.children))
self.assertEqual('simple_stmt',
pytree_utils.NodeName(tree.children[0]))
def testInsertNodesAfter(self):
# Insert after and make sure it went to the right place
pytree_utils.InsertNodesAfter([self._MakeNewNodeRPAR()],
self._simple_tree.children[2])
self.assertEqual(4, len(self._simple_tree.children))
self.assertEqual('simple_stmt',
pytree_utils.NodeName(self._simple_tree.children[2]))
self.assertEqual('RPAR',
pytree_utils.NodeName(self._simple_tree.children[3]))
def testInsertNodesAfterLastChild(self):
# Insert after the last child of its parent
simple_stmt = self._simple_tree.children[2]
foo_child = simple_stmt.children[0]
pytree_utils.InsertNodesAfter([self._MakeNewNodeRPAR()], foo_child)
self.assertEqual(3, len(self._simple_tree.children))
self.assertEqual(2, len(simple_stmt.children))
self.assertEqual('NAME',
pytree_utils.NodeName(simple_stmt.children[0]))
self.assertEqual('RPAR',
pytree_utils.NodeName(simple_stmt.children[1]))
def _FindNthChildNamed(self, node, name, n=1):
for i, child in enumerate(
py3compat.ifilter(lambda c: pytree_utils.NodeName(c) == name,
node.pre_order())):
if i == n - 1:
return child
raise RuntimeError('No Nth child for n={0}'.format(n))
def _AssertNodeIsComment(self, node, text_in_comment=None):
if pytree_utils.NodeName(node) == 'simple_stmt':
self._AssertNodeType('COMMENT', node.children[0])
node_value = node.children[0].value
else:
self._AssertNodeType('COMMENT', node)
node_value = node.value
if text_in_comment is not None:
self.assertIn(text_in_comment, node_value)
def _SetExpressionOperandPenalty(node, ops):
for index in py3compat.range(1, len(node.children) - 1):
child = node.children[index]
if pytree_utils.NodeName(child) in ops:
if style.Get('SPLIT_BEFORE_ARITHMETIC_OPERATOR'):
_SetSplitPenalty(
child, style.Get('SPLIT_PENALTY_ARITHMETIC_OPERATOR'))
else:
_SetSplitPenalty(
pytree_utils.FirstLeafNode(node.children[index + 1]),
style.Get('SPLIT_PENALTY_ARITHMETIC_OPERATOR'))
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 Visit(self, node):
"""Visit a node."""
method = 'Visit_{0}'.format(pytree_utils.NodeName(node))
if hasattr(self, method):
# Found a specific visitor for this node
getattr(self, method)(node)
else:
if isinstance(node, pytree.Leaf):
self.DefaultLeafVisit(node)
else:
self.DefaultNodeVisit(node)
def Visit_comparison(self, node): # pylint: disable=invalid-name
# comparison ::= expr (comp_op expr)*
# comp_op ::= '<'|'>'|'=='|'>='|'<='|'<>'|'!='|'in'|'not in'|'is'|'is not'
for child in node.children:
self.Visit(child)
if (isinstance(child, pytree.Leaf) and child.value
in {'<', '>', '==', '>=', '<=', '<>', '!=', 'in', 'is'}):
_AppendTokenSubtype(child, subtypes.BINARY_OPERATOR)
elif pytree_utils.NodeName(child) == 'comp_op':
for grandchild in child.children:
_AppendTokenSubtype(grandchild, subtypes.BINARY_OPERATOR)
def _StronglyConnectedCompOp(op):
if (len(op.children[1].children) == 2
and pytree_utils.NodeName(op.children[1]) == 'comp_op'):
if (pytree_utils.FirstLeafNode(op.children[1]).value == 'not'
and pytree_utils.LastLeafNode(op.children[1]).value == 'in'):
return True
if (pytree_utils.FirstLeafNode(op.children[1]).value == 'is'
and pytree_utils.LastLeafNode(op.children[1]).value == 'not'):
return True
if (isinstance(op.children[1], pytree.Leaf)
and op.children[1].value in _COMP_OPS):
return True
return False
def _FindStmtParent(node):
"""Find the nearest parent of node that is a statement node.
Arguments:
node: node to start from
Returns:
Nearest parent (or node itself, if suitable).
"""
if pytree_utils.NodeName(node) in _STATEMENT_NODES:
return node
else:
return _FindStmtParent(node.parent)
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 Visit_typedargslist(self, node): # pylint: disable=invalid-name
# typedargslist ::=
# ((tfpdef ['=' test] ',')*
# ('*' [tname] (',' tname ['=' test])* [',' '**' tname]
# | '**' tname)
# | tfpdef ['=' test] (',' tfpdef ['=' test])* [','])
self._ProcessArgLists(node)
_SetArgListSubtype(node, subtypes.DEFAULT_OR_NAMED_ASSIGN,
subtypes.DEFAULT_OR_NAMED_ASSIGN_ARG_LIST)
tname = False
if not node.children:
return
_AppendFirstLeafTokenSubtype(node.children[0],
subtypes.PARAMETER_START)
_AppendLastLeafTokenSubtype(node.children[-1], subtypes.PARAMETER_STOP)
tname = pytree_utils.NodeName(node.children[0]) == 'tname'
for i in range(1, len(node.children)):
prev_child = node.children[i - 1]
child = node.children[i]
if prev_child.type == grammar_token.COMMA:
_AppendFirstLeafTokenSubtype(child, subtypes.PARAMETER_START)
elif child.type == grammar_token.COMMA:
_AppendLastLeafTokenSubtype(prev_child,
subtypes.PARAMETER_STOP)
if pytree_utils.NodeName(child) == 'tname':
tname = True
_SetArgListSubtype(child, subtypes.TYPED_NAME,
subtypes.TYPED_NAME_ARG_LIST)
elif child.type == grammar_token.COMMA:
tname = False
elif child.type == grammar_token.EQUAL and tname:
_AppendTokenSubtype(child, subtype=subtypes.TYPED_NAME)
tname = False
def DefaultNodeVisit(self, node):
"""Override the default visitor for Node.
This will set the blank lines required if the last entity was a class or
function.
Arguments:
node: (pytree.Node) The node to visit.
"""
if self.last_was_class_or_function:
if pytree_utils.NodeName(node) in _PYTHON_STATEMENTS:
leaf = pytree_utils.FirstLeafNode(node)
_SetNumNewlines(leaf, self._GetNumNewlines(leaf))
self.last_was_class_or_function = False
super(_BlankLineCalculator, self).DefaultNodeVisit(node)
def Visit_dictsetmaker(self, node): # pylint: disable=invalid-name
# dictsetmaker ::= (test ':' test (comp_for |
# (',' test ':' test)* [','])) |
# (test (comp_for | (',' test)* [',']))
for child in node.children:
self.Visit(child)
comp_for = False
dict_maker = False
for child in node.children:
if pytree_utils.NodeName(child) == 'comp_for':
comp_for = True
_AppendFirstLeafTokenSubtype(child,
subtypes.DICT_SET_GENERATOR)
elif child.type in (grammar_token.COLON, grammar_token.DOUBLESTAR):
dict_maker = True
if not comp_for and dict_maker:
last_was_colon = False
unpacking = False
for child in node.children:
if child.type == grammar_token.DOUBLESTAR:
_AppendFirstLeafTokenSubtype(child,
subtypes.KWARGS_STAR_STAR)
if last_was_colon:
if style.Get('INDENT_DICTIONARY_VALUE'):
_InsertPseudoParentheses(child)
else:
_AppendFirstLeafTokenSubtype(child,
subtypes.DICTIONARY_VALUE)
elif (isinstance(child, pytree.Node)
or (not child.value.startswith('#')
and child.value not in '{:,')):
# Mark the first leaf of a key entry as a DICTIONARY_KEY. We
# normally want to split before them if the dictionary cannot exist
# on a single line.
if not unpacking or pytree_utils.FirstLeafNode(
child).value == '**':
_AppendFirstLeafTokenSubtype(child,
subtypes.DICTIONARY_KEY)
_AppendSubtypeRec(child, subtypes.DICTIONARY_KEY_PART)
last_was_colon = child.type == grammar_token.COLON
if child.type == grammar_token.DOUBLESTAR:
unpacking = True
elif last_was_colon:
unpacking = False
def _FindNodeWithStandaloneLineParent(node):
"""Find a node whose parent is a 'standalone line' node.
See the comment above _STANDALONE_LINE_NODES for more details.
Arguments:
node: node to start from
Returns:
Suitable node that's either the node itself or one of its ancestors.
"""
if pytree_utils.NodeName(node.parent) in _STANDALONE_LINE_NODES:
return node
else:
# This is guaranteed to terminate because 'file_input' is the root node of
# any pytree.
return _FindNodeWithStandaloneLineParent(node.parent)
def Visit_atom(self, node): # pylint: disable=invalid-name
# atom ::= ('(' [yield_expr|testlist_gexp] ')'
# '[' [listmaker] ']' |
# '{' [dictsetmaker] '}')
self.DefaultNodeVisit(node)
if (node.children[0].value == '('
and not hasattr(node.children[0], 'is_pseudo')):
if node.children[-1].value == ')':
if pytree_utils.NodeName(node.parent) == 'if_stmt':
_SetSplitPenalty(node.children[-1], STRONGLY_CONNECTED)
else:
if len(node.children) > 2:
_SetSplitPenalty(
pytree_utils.FirstLeafNode(node.children[1]), EXPR)
_SetSplitPenalty(node.children[-1], ATOM)
elif node.children[0].value in '[{' and len(node.children) == 2:
# Keep empty containers together if we can.
_SetUnbreakable(node.children[-1])
def Visit_arglist(self, node): # pylint: disable=invalid-name
# arglist ::= argument (',' argument)* [',']
if node.children[0].type == grammar_token.STAR:
# Python 3 treats a star expression as a specific expression type.
# Process it in that method.
self.Visit_star_expr(node)
return
self.DefaultNodeVisit(node)
for index in py3compat.range(1, len(node.children)):
child = node.children[index]
if isinstance(child, pytree.Leaf) and child.value == ',':
_SetUnbreakable(child)
for child in node.children:
if pytree_utils.NodeName(child) == 'atom':
_IncreasePenalty(child, CONNECTED)
def Visit_simple_stmt(self, node):
# A 'simple_stmt' conveniently represents a non-compound Python statement,
# i.e. a statement that does not contain other statements.
# When compound nodes have a single statement as their suite, the parser
# can leave it in the tree directly without creating a suite. But we have
# to increase depth in these cases as well. However, don't increase the
# depth of we have a simple_stmt that's a comment node. This represents a
# standalone comment and in the case of it coming directly after the
# funcdef, it is a "top" comment for the whole function.
# TODO(eliben): add more relevant compound statements here.
single_stmt_suite = (node.parent and pytree_utils.NodeName(node.parent)
in self._STMT_TYPES)
is_comment_stmt = pytree_utils.IsCommentStatement(node)
if single_stmt_suite and not is_comment_stmt:
self._cur_depth += 1
self._StartNewLine()
self.DefaultNodeVisit(node)
if single_stmt_suite and not is_comment_stmt:
self._cur_depth -= 1
def __init__(self, node, node_name=None):
"""Constructor.
Arguments:
node: (pytree.Leaf) The node that's being wrapped.
none_name: (string) The name of the node.
"""
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) if not node_name else node_name
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 _SetArgListSubtype(node, node_subtype, list_subtype):
"""Set named assign subtype on elements in a arg list."""
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
if not HasSubtype(node):
return
for child in node.children:
node_name = pytree_utils.NodeName(child)
if node_name not in {'atom', 'COMMA'}:
_AppendFirstLeafTokenSubtype(child, list_subtype)
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
_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
_SetUnbreakable(node.children[colon_idx])
self.DefaultNodeVisit(node)
if arrow_idx > 0:
_SetSplitPenalty(
pytree_utils.LastLeafNode(node.children[arrow_idx - 1]), 0)
_SetUnbreakable(node.children[arrow_idx])
_SetStronglyConnected(node.children[arrow_idx + 1])
def _HasPrecedence(tok):
"""Whether a binary operation has precedence within its context."""
node = tok.node
# We let ancestor be the statement surrounding the operation that tok is the
# operator in.
ancestor = node.parent.parent
while ancestor is not None:
# Search through the ancestor nodes in the parse tree for operators with
# lower precedence.
predecessor_type = pytree_utils.NodeName(ancestor)
if predecessor_type in ['arith_expr', 'term']:
# An ancestor "arith_expr" or "term" means we have found an operator
# with lower precedence than our tok.
return True
if predecessor_type != 'atom':
# We understand the context to look for precedence within as an
# arbitrary nesting of "arith_expr", "term", and "atom" nodes. If we
# leave this context we have not found a lower precedence operator.
return False
# Under normal usage we expect a complete parse tree to be available and
# we will return before we get an AttributeError from the root.
ancestor = ancestor.parent
def testNodeNameForNode(self):
leaf = pytree.Leaf(token.LPAR, '(')
node = pytree.Node(pygram.python_grammar.symbol2number['suite'],
[leaf])
self.assertEqual('suite', pytree_utils.NodeName(node))
def testPrintStatementToTree(self):
tree = pytree_utils.ParseCodeToTree('print "hello world"\n')
self.assertEqual('file_input', pytree_utils.NodeName(tree))
self.assertEqual(2, len(tree.children))
self.assertEqual('simple_stmt',
pytree_utils.NodeName(tree.children[0]))
def testClassNotLocal(self):
tree = pytree_utils.ParseCodeToTree('class nonlocal: pass\n')
self.assertEqual('file_input', pytree_utils.NodeName(tree))
self.assertEqual(2, len(tree.children))
self.assertEqual('classdef', pytree_utils.NodeName(tree.children[0]))
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 _VisitNodeRec(node):
"""Recursively visit each node to splice comments into the AST."""
# This loop may insert into node.children, so we'll iterate over a copy.
for child in node.children[:]:
if isinstance(child, pytree.Node):
# Nodes don't have prefixes.
_VisitNodeRec(child)
else:
if child.prefix.lstrip().startswith('#'):
# We have a comment prefix in this child, so splicing is needed.
comment_prefix = child.prefix
comment_lineno = child.lineno - comment_prefix.count('\n')
comment_column = child.column
# Remember the leading indentation of this prefix and clear it.
# Mopping up the prefix is important because we may go over this same
# child in the next iteration...
child_prefix = child.prefix.lstrip('\n')
prefix_indent = child_prefix[:child_prefix.find('#')]
if '\n' in prefix_indent:
prefix_indent = prefix_indent[prefix_indent.rfind('\n') + 1:]
child.prefix = ''
if child.type == token.NEWLINE:
# If the prefix was on a NEWLINE leaf, it's part of the line so it
# will be inserted after the previously encountered leaf.
# We can't just insert it before the NEWLINE node, because as a
# result of the way pytrees are organized, this node can be under
# an inappropriate parent.
comment_column -= len(comment_prefix.lstrip())
pytree_utils.InsertNodesAfter(
_CreateCommentsFromPrefix(
comment_prefix,
comment_lineno,
comment_column,
standalone=False), prev_leaf[0])
elif child.type == token.DEDENT:
# Comment prefixes on DEDENT nodes also deserve special treatment,
# because their final placement depends on their prefix.
# We'll look for an ancestor of this child with a matching
# indentation, and insert the comment before it if the ancestor is
# on a DEDENT node and after it otherwise.
#
# lib2to3 places comments that should be separated into the same
# DEDENT node. For example, "comment 1" and "comment 2" will be
# combined.
#
# def _():
# for x in y:
# pass
# # comment 1
#
# # comment 2
# pass
#
# In this case, we need to split them up ourselves.
# Split into groups of comments at decreasing levels of indentation
comment_groups = []
comment_column = None
for cmt in comment_prefix.split('\n'):
col = cmt.find('#')
if col < 0:
if comment_column is None:
# Skip empty lines at the top of the first comment group
comment_lineno += 1
continue
elif comment_column is None or col < comment_column:
comment_column = col
comment_indent = cmt[:comment_column]
comment_groups.append((comment_column, comment_indent, []))
comment_groups[-1][-1].append(cmt)
# Insert a node for each group
for comment_column, comment_indent, comment_group in comment_groups:
ancestor_at_indent = _FindAncestorAtIndent(child, comment_indent)
if ancestor_at_indent.type == token.DEDENT:
InsertNodes = pytree_utils.InsertNodesBefore # pylint: disable=invalid-name # noqa
else:
InsertNodes = pytree_utils.InsertNodesAfter # pylint: disable=invalid-name # noqa
InsertNodes(
_CreateCommentsFromPrefix(
'\n'.join(comment_group) + '\n',
comment_lineno,
comment_column,
standalone=True), ancestor_at_indent)
comment_lineno += len(comment_group)
else:
# Otherwise there are two cases.
#
# 1. The comment is on its own line
# 2. The comment is part of an expression.
#
# Unfortunately, it's fairly difficult to distinguish between the
# two in lib2to3 trees. The algorithm here is to determine whether
# child is the first leaf in the statement it belongs to. If it is,
# then the comment (which is a prefix) belongs on a separate line.
# If it is not, it means the comment is buried deep in the statement
# and is part of some expression.
stmt_parent = _FindStmtParent(child)
for leaf_in_parent in stmt_parent.leaves():
if leaf_in_parent.type == token.NEWLINE:
continue
elif id(leaf_in_parent) == id(child):
# This comment stands on its own line, and it has to be inserted
# into the appropriate parent. We'll have to find a suitable
# parent to insert into. See comments above
# _STANDALONE_LINE_NODES for more details.
node_with_line_parent = _FindNodeWithStandaloneLineParent(child)
if pytree_utils.NodeName(
node_with_line_parent.parent) in {'funcdef', 'classdef'}:
# Keep a comment that's not attached to a function or class
# next to the object it is attached to.
comment_end = (
comment_lineno + comment_prefix.rstrip('\n').count('\n'))
if comment_end < node_with_line_parent.lineno - 1:
node_with_line_parent = node_with_line_parent.parent
pytree_utils.InsertNodesBefore(
_CreateCommentsFromPrefix(
comment_prefix, comment_lineno, 0, standalone=True),
node_with_line_parent)
break
else:
if comment_lineno == prev_leaf[0].lineno:
comment_lines = comment_prefix.splitlines()
value = comment_lines[0].lstrip()
if value.rstrip('\n'):
comment_column = prev_leaf[0].column
comment_column += len(prev_leaf[0].value)
comment_column += (
len(comment_lines[0]) - len(comment_lines[0].lstrip()))
comment_leaf = pytree.Leaf(
type=token.COMMENT,
value=value.rstrip('\n'),
context=('', (comment_lineno, comment_column)))
pytree_utils.InsertNodesAfter([comment_leaf], prev_leaf[0])
comment_prefix = '\n'.join(comment_lines[1:])
comment_lineno += 1
rindex = (0 if '\n' not in comment_prefix.rstrip() else
comment_prefix.rstrip().rindex('\n') + 1)
comment_column = (
len(comment_prefix[rindex:]) -
len(comment_prefix[rindex:].lstrip()))
comments = _CreateCommentsFromPrefix(
comment_prefix,
comment_lineno,
comment_column,
standalone=False)
pytree_utils.InsertNodesBefore(comments, child)
break
prev_leaf[0] = child
