| # Copyright 2015-2017 Google Inc. All Rights Reserved. |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); |
| # you may not use this file except in compliance with the License. |
| # You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| """Comment splicer for lib2to3 trees. |
| |
| The lib2to3 syntax tree produced by the parser holds comments and whitespace in |
| prefix attributes of nodes, rather than nodes themselves. This module provides |
| functionality to splice comments out of prefixes and into nodes of their own, |
| making them easier to process. |
| |
| SpliceComments(): the main function exported by this module. |
| """ |
| |
| from lib2to3 import pygram |
| from lib2to3 import pytree |
| from lib2to3.pgen2 import token |
| |
| from yapf.yapflib import pytree_utils |
| |
| |
| def SpliceComments(tree): |
| """Given a pytree, splice comments into nodes of their own right. |
| |
| Extract comments from the prefixes where they are housed after parsing. |
| The prefixes that previously housed the comments become empty. |
| |
| Args: |
| tree: a pytree.Node - the tree to work on. The tree is modified by this |
| function. |
| """ |
| # The previous leaf node encountered in the traversal. |
| # This is a list because Python 2.x doesn't have 'nonlocal' :) |
| prev_leaf = [None] |
| _AnnotateIndents(tree) |
| |
| 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 |
| |
| _VisitNodeRec(tree) |
| |
| |
| def _CreateCommentsFromPrefix(comment_prefix, |
| comment_lineno, |
| comment_column, |
| standalone=False): |
| """Create pytree nodes to represent the given comment prefix. |
| |
| Args: |
| comment_prefix: (unicode) the text of the comment from the node's prefix. |
| comment_lineno: (int) the line number for the start of the comment. |
| comment_column: (int) the column for the start of the comment. |
| standalone: (bool) determines if the comment is standalone or not. |
| |
| Returns: |
| The simple_stmt nodes if this is a standalone comment, otherwise a list of |
| new COMMENT leafs. The prefix may consist of multiple comment blocks, |
| separated by blank lines. Each block gets its own leaf. |
| """ |
| # The comment is stored in the prefix attribute, with no lineno of its |
| # own. So we only know at which line it ends. To find out at which line it |
| # starts, look at how many newlines the comment itself contains. |
| comments = [] |
| |
| lines = comment_prefix.split('\n') |
| index = 0 |
| while index < len(lines): |
| comment_block = [] |
| while index < len(lines) and lines[index].lstrip().startswith('#'): |
| comment_block.append(lines[index].strip()) |
| index += 1 |
| |
| if comment_block: |
| new_lineno = comment_lineno + index - 1 |
| comment_block[0] = comment_block[0].strip() |
| comment_block[-1] = comment_block[-1].strip() |
| comment_leaf = pytree.Leaf( |
| type=token.COMMENT, |
| value='\n'.join(comment_block), |
| context=('', (new_lineno, comment_column))) |
| comment_node = comment_leaf if not standalone else pytree.Node( |
| pygram.python_symbols.simple_stmt, [comment_leaf]) |
| comments.append(comment_node) |
| |
| while index < len(lines) and not lines[index].lstrip(): |
| index += 1 |
| |
| return comments |
| |
| |
| # "Standalone line nodes" are tree nodes that have to start a new line in Python |
| # code (and cannot follow a ';' or ':'). Other nodes, like 'expr_stmt', serve as |
| # parents of other nodes but can come later in a line. This is a list of |
| # standalone line nodes in the grammar. It is meant to be exhaustive |
| # *eventually*, and we'll modify it with time as we discover more corner cases |
| # in the parse tree. |
| # |
| # When splicing a standalone comment (i.e. a comment that appears on its own |
| # line, not on the same line with other code), it's important to insert it into |
| # an appropriate parent of the node it's attached to. An appropriate parent |
| # is the first "standaline line node" in the parent chain of a node. |
| _STANDALONE_LINE_NODES = frozenset([ |
| 'suite', 'if_stmt', 'while_stmt', 'for_stmt', 'try_stmt', 'with_stmt', |
| 'funcdef', 'classdef', 'decorated', 'file_input' |
| ]) |
| |
| |
| 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) |
| |
| |
| # "Statement nodes" are standalone statements. The don't have to start a new |
| # line. |
| _STATEMENT_NODES = frozenset(['simple_stmt']) | _STANDALONE_LINE_NODES |
| |
| |
| 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 _FindAncestorAtIndent(node, indent): |
| """Find an ancestor of node with the given indentation. |
| |
| Arguments: |
| node: node to start from. This must not be the tree root. |
| indent: indentation string for the ancestor we're looking for. |
| See _AnnotateIndents for more details. |
| |
| Returns: |
| An ancestor node with suitable indentation. If no suitable ancestor is |
| found, the closest ancestor to the tree root is returned. |
| """ |
| if node.parent.parent is None: |
| # Our parent is the tree root, so there's nowhere else to go. |
| return node |
| |
| # If the parent has an indent annotation, and it's shorter than node's |
| # indent, this is a suitable ancestor. |
| # The reason for "shorter" rather than "equal" is that comments may be |
| # improperly indented (i.e. by three spaces, where surrounding statements |
| # have either zero or two or four), and we don't want to propagate them all |
| # the way to the root. |
| parent_indent = pytree_utils.GetNodeAnnotation( |
| node.parent, pytree_utils.Annotation.CHILD_INDENT) |
| if parent_indent is not None and indent.startswith(parent_indent): |
| return node |
| else: |
| # Keep looking up the tree. |
| return _FindAncestorAtIndent(node.parent, indent) |
| |
| |
| def _FindNextAncestor(node): |
| if node.parent is None: |
| return node |
| |
| if node.parent.next_sibling is not None: |
| return node.parent.next_sibling |
| |
| return _FindNextAncestor(node.parent) |
| |
| |
| 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) |