diff options
Diffstat (limited to 'yapf/yapflib/pytree_unwrapper.py')
-rw-r--r-- | yapf/yapflib/pytree_unwrapper.py | 117 |
1 files changed, 77 insertions, 40 deletions
diff --git a/yapf/yapflib/pytree_unwrapper.py b/yapf/yapflib/pytree_unwrapper.py index 0d371ae..1b05b0e 100644 --- a/yapf/yapflib/pytree_unwrapper.py +++ b/yapf/yapflib/pytree_unwrapper.py @@ -11,13 +11,13 @@ # 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. -"""PyTreeUnwrapper - produces a list of unwrapped lines from a pytree. +"""PyTreeUnwrapper - produces a list of logical lines from a pytree. -[for a description of what an unwrapped line is, see unwrapped_line.py] +[for a description of what a logical line is, see logical_line.py] This is a pytree visitor that goes over a parse tree and produces a list of -UnwrappedLine containers from it, each with its own depth and containing all -the tokens that could fit on the line if there were no maximal line-length +LogicalLine containers from it, each with its own depth and containing all the +tokens that could fit on the line if there were no maximal line-length limitations. Note: a precondition to running this visitor and obtaining correct results is @@ -31,27 +31,30 @@ For most uses, the convenience function UnwrapPyTree should be sufficient. from lib2to3 import pytree from lib2to3.pgen2 import token as grammar_token +from yapf.yapflib import format_token +from yapf.yapflib import logical_line +from yapf.yapflib import object_state from yapf.yapflib import pytree_utils from yapf.yapflib import pytree_visitor from yapf.yapflib import split_penalty from yapf.yapflib import style -from yapf.yapflib import unwrapped_line +from yapf.yapflib import subtypes def UnwrapPyTree(tree): - """Create and return a list of unwrapped lines from the given pytree. + """Create and return a list of logical lines from the given pytree. Arguments: - tree: the top-level pytree node to unwrap. + tree: the top-level pytree node to unwrap.. Returns: - A list of UnwrappedLine objects. + A list of LogicalLine objects. """ unwrapper = PyTreeUnwrapper() unwrapper.Visit(tree) - uwlines = unwrapper.GetUnwrappedLines() - uwlines.sort(key=lambda x: x.lineno) - return uwlines + llines = unwrapper.GetLogicalLines() + llines.sort(key=lambda x: x.lineno) + return llines # Grammar tokens considered as whitespace for the purpose of unwrapping. @@ -77,39 +80,40 @@ class PyTreeUnwrapper(pytree_visitor.PyTreeVisitor): """ def __init__(self): - # A list of all unwrapped lines finished visiting so far. - self._unwrapped_lines = [] + # A list of all logical lines finished visiting so far. + self._logical_lines = [] - # Builds up a "current" unwrapped line while visiting pytree nodes. Some - # nodes will finish a line and start a new one. - self._cur_unwrapped_line = unwrapped_line.UnwrappedLine(0) + # Builds up a "current" logical line while visiting pytree nodes. Some nodes + # will finish a line and start a new one. + self._cur_logical_line = logical_line.LogicalLine(0) # Current indentation depth. self._cur_depth = 0 - def GetUnwrappedLines(self): + def GetLogicalLines(self): """Fetch the result of the tree walk. Note: only call this after visiting the whole tree. Returns: - A list of UnwrappedLine objects. + A list of LogicalLine objects. """ # Make sure the last line that was being populated is flushed. self._StartNewLine() - return self._unwrapped_lines + return self._logical_lines def _StartNewLine(self): """Finish current line and start a new one. - Place the currently accumulated line into the _unwrapped_lines list and + Place the currently accumulated line into the _logical_lines list and start a new one. """ - if self._cur_unwrapped_line.tokens: - self._unwrapped_lines.append(self._cur_unwrapped_line) - _MatchBrackets(self._cur_unwrapped_line) - _AdjustSplitPenalty(self._cur_unwrapped_line) - self._cur_unwrapped_line = unwrapped_line.UnwrappedLine(self._cur_depth) + if self._cur_logical_line.tokens: + self._logical_lines.append(self._cur_logical_line) + _MatchBrackets(self._cur_logical_line) + _IdentifyParameterLists(self._cur_logical_line) + _AdjustSplitPenalty(self._cur_logical_line) + self._cur_logical_line = logical_line.LogicalLine(self._cur_depth) _STMT_TYPES = frozenset({ 'if_stmt', @@ -148,7 +152,7 @@ class PyTreeUnwrapper(pytree_visitor.PyTreeVisitor): """Helper for visiting compound statements. Python compound statements serve as containers for other statements. Thus, - when we encounter a new compound statement we start a new unwrapped line. + when we encounter a new compound statement, we start a new logical line. Arguments: node: the node to visit. @@ -202,7 +206,7 @@ class PyTreeUnwrapper(pytree_visitor.PyTreeVisitor): for child in node.children: index += 1 self.Visit(child) - if pytree_utils.NodeName(child) == 'ASYNC': + if child.type == grammar_token.ASYNC: break for child in node.children[index].children: self.Visit(child) @@ -218,16 +222,17 @@ class PyTreeUnwrapper(pytree_visitor.PyTreeVisitor): for child in node.children: index += 1 self.Visit(child) - if pytree_utils.NodeName(child) == 'ASYNC': + if child.type == grammar_token.ASYNC: break for child in node.children[index].children: + if child.type == grammar_token.NAME and child.value == 'else': + self._StartNewLine() self.Visit(child) def Visit_decorator(self, node): # pylint: disable=invalid-name for child in node.children: self.Visit(child) - if (pytree_utils.NodeName(child) == 'COMMENT' and - child == node.children[0]): + if child.type == grammar_token.COMMENT and child == node.children[0]: self._StartNewLine() def Visit_decorators(self, node): # pylint: disable=invalid-name @@ -280,7 +285,7 @@ class PyTreeUnwrapper(pytree_visitor.PyTreeVisitor): def DefaultLeafVisit(self, leaf): """Default visitor for tree leaves. - A tree leaf is always just gets appended to the current unwrapped line. + A tree leaf is always just gets appended to the current logical line. Arguments: leaf: the leaf to visit. @@ -289,13 +294,13 @@ class PyTreeUnwrapper(pytree_visitor.PyTreeVisitor): self._StartNewLine() elif leaf.type != grammar_token.COMMENT or leaf.value.strip(): # Add non-whitespace tokens and comments that aren't empty. - self._cur_unwrapped_line.AppendNode(leaf) + self._cur_logical_line.AppendNode(leaf) _BRACKET_MATCH = {')': '(', '}': '{', ']': '['} -def _MatchBrackets(uwline): +def _MatchBrackets(line): """Visit the node and match the brackets. For every open bracket ('[', '{', or '('), find the associated closing bracket @@ -303,10 +308,10 @@ def _MatchBrackets(uwline): or close bracket. Arguments: - uwline: (UnwrappedLine) An unwrapped line. + line: (LogicalLine) A logical line. """ bracket_stack = [] - for token in uwline.tokens: + for token in line.tokens: if token.value in pytree_utils.OPENING_BRACKETS: bracket_stack.append(token) elif token.value in pytree_utils.CLOSING_BRACKETS: @@ -320,17 +325,50 @@ def _MatchBrackets(uwline): token.container_opening = bracket -def _AdjustSplitPenalty(uwline): +def _IdentifyParameterLists(line): + """Visit the node to create a state for parameter lists. + + For instance, a parameter is considered an "object" with its first and last + token uniquely identifying the object. + + Arguments: + line: (LogicalLine) A logical line. + """ + func_stack = [] + param_stack = [] + for tok in line.tokens: + # Identify parameter list objects. + if subtypes.FUNC_DEF in tok.subtypes: + assert tok.next_token.value == '(' + func_stack.append(tok.next_token) + continue + + if func_stack and tok.value == ')': + if tok == func_stack[-1].matching_bracket: + func_stack.pop() + continue + + # Identify parameter objects. + if subtypes.PARAMETER_START in tok.subtypes: + param_stack.append(tok) + + # Not "elif", a parameter could be a single token. + if param_stack and subtypes.PARAMETER_STOP in tok.subtypes: + start = param_stack.pop() + func_stack[-1].parameters.append(object_state.Parameter(start, tok)) + + +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: - uwline: (UnwrappedLine) An unwrapped line. + line: (LogicalLine) An logical line. """ bracket_level = 0 - for index, token in enumerate(uwline.tokens): + for index, token in enumerate(line.tokens): if index and not bracket_level: pytree_utils.SetNodeAnnotation(token.node, pytree_utils.Annotation.SPLIT_PENALTY, @@ -348,8 +386,7 @@ def _DetermineMustSplitAnnotation(node): if not _ContainsComments(node): token = next(node.parent.leaves()) if token.value == '(': - if sum(1 for ch in node.children - if pytree_utils.NodeName(ch) == 'COMMA') < 2: + if sum(1 for ch in node.children if ch.type == grammar_token.COMMA) < 2: return if (not isinstance(node.children[-1], pytree.Leaf) or node.children[-1].value != ','): |