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 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 _VisitNodeRec(node):
# 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)
comment_column += len(comment_prefix) - 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 after it.
ancestor_at_indent = _FindAncestorAtIndent(child, prefix_indent)
if ancestor_at_indent.type == token.DEDENT:
comments = comment_prefix.split('\n')
# 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.
before = []
after = []
after_lineno = comment_lineno
index = 0
while index < len(comments):
cmt = comments[index]
if not cmt.strip() or cmt.startswith(prefix_indent + '#'):
before.append(cmt)
else:
after_lineno += index
after.extend(comments[index:])
break
index += 1
# Special case where the comment is inserted in the same
# indentation level as the DEDENT it was originally attached to.
pytree_utils.InsertNodesBefore(
_CreateCommentsFromPrefix(
'\n'.join(before) + '\n',
comment_lineno,
comment_column,
standalone=True), ancestor_at_indent)
if after:
after_column = len(after[0]) - len(after[0].lstrip())
comment_column -= comment_column - after_column
pytree_utils.InsertNodesAfter(
_CreateCommentsFromPrefix(
'\n'.join(after) + '\n',
after_lineno,
comment_column,
standalone=True), _FindNextAncestor(ancestor_at_indent))
else:
pytree_utils.InsertNodesAfter(
_CreateCommentsFromPrefix(
comment_prefix,
comment_lineno,
comment_column,
standalone=True), ancestor_at_indent)
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)
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
def _VisitNodeRec(node):
# 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
else:
InsertNodes = pytree_utils.InsertNodesAfter # pylint: disable=invalid-name
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
def testInsertNodesWhichHasParent(self):
# Try to insert an existing tree node into another place and fail.
with self.assertRaises(RuntimeError):
pytree_utils.InsertNodesAfter([self._simple_tree.children[1]],
self._simple_tree.children[0])
def _VisitNodeRec(node):
# 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')
# 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('#')]
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.
assert prev_leaf[0] is not None
pytree_utils.InsertNodesAfter(
_CreateCommentsFromPrefix(comment_prefix,
comment_lineno,
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 after it.
ancestor_at_indent = _FindAncestorAtIndent(
child, prefix_indent)
if ancestor_at_indent.type == token.DEDENT:
# Special case where the comment is inserted in the same
# indentation level as the DEDENT it was originally attached to.
pytree_utils.InsertNodesBefore(
_CreateCommentsFromPrefix(comment_prefix,
comment_lineno,
standalone=True),
ancestor_at_indent)
else:
pytree_utils.InsertNodesAfter(
_CreateCommentsFromPrefix(comment_prefix,
comment_lineno,
standalone=True),
ancestor_at_indent)
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)
pytree_utils.InsertNodesBefore(
_CreateCommentsFromPrefix(comment_prefix,
comment_lineno,
standalone=True),
node_with_line_parent)
break
else:
if comment_lineno == prev_leaf[0].lineno:
comment_lines = comment_prefix.splitlines()
comment_leaf = pytree.Leaf(
type=token.COMMENT,
value=comment_lines[0].strip(),
context=('', (comment_lineno, 0)))
pytree_utils.InsertNodesAfter(
[comment_leaf], prev_leaf[0])
comment_prefix = '\n'.join(
comment_lines[1:])
comment_lineno += 1
comments = _CreateCommentsFromPrefix(
comment_prefix,
comment_lineno,
standalone=False)
pytree_utils.InsertNodesBefore(comments, child)
break
prev_leaf[0] = child
