| # Copyright 2015-2016 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. |
| """Computation of split penalties before/between tokens.""" |
| |
| from lib2to3 import pytree |
| |
| from yapf.yapflib import format_token |
| from yapf.yapflib import py3compat |
| from yapf.yapflib import pytree_utils |
| from yapf.yapflib import pytree_visitor |
| from yapf.yapflib import style |
| |
| # TODO(morbo): Document the annotations in a centralized place. E.g., the |
| # README file. |
| UNBREAKABLE = 1000 * 1000 |
| VERY_STRONGLY_CONNECTED = 3000 |
| DOTTED_NAME = 2500 |
| STRONGLY_CONNECTED = 2000 |
| CONTIGUOUS_LIST = 500 |
| |
| NOT_TEST = 242 |
| COMPARISON_EXPRESSION = 842 |
| ARITHMETIC_EXPRESSION = 942 |
| |
| |
| def ComputeSplitPenalties(tree): |
| """Compute split penalties on tokens in the given parse tree. |
| |
| Arguments: |
| tree: the top-level pytree node to annotate with penalties. |
| """ |
| _TreePenaltyAssigner().Visit(tree) |
| |
| |
| class _TreePenaltyAssigner(pytree_visitor.PyTreeVisitor): |
| """Assigns split penalties to tokens, based on parse tree structure. |
| |
| Split penalties are attached as annotations to tokens. |
| """ |
| |
| 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_classdef(self, node): # pylint: disable=invalid-name |
| # classdef ::= 'class' NAME ['(' [arglist] ')'] ':' suite |
| # |
| # NAME |
| self._SetUnbreakable(node.children[1]) |
| if len(node.children) > 4: |
| # opening '(' |
| self._SetUnbreakable(node.children[2]) |
| # ':' |
| self._SetUnbreakable(node.children[-2]) |
| self.DefaultNodeVisit(node) |
| |
| 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]) |
| while colon_idx < len(node.children): |
| if (isinstance(node.children[colon_idx], pytree.Leaf) and |
| node.children[colon_idx].value == ':'): |
| break |
| colon_idx += 1 |
| self._SetUnbreakable(node.children[colon_idx]) |
| self.DefaultNodeVisit(node) |
| |
| def Visit_lambdef(self, node): # pylint: disable=invalid-name |
| # lambdef ::= 'lambda' [varargslist] ':' test |
| # Loop over the lambda up to and including the colon. |
| if style.Get('ALLOW_MULTILINE_LAMBDAS'): |
| self._SetStronglyConnected(node) |
| else: |
| self._SetUnbreakableOnChildren(node) |
| |
| def Visit_parameters(self, node): # pylint: disable=invalid-name |
| # parameters ::= '(' [typedargslist] ')' |
| self.DefaultNodeVisit(node) |
| |
| # Can't break before the opening paren of a parameter list. |
| self._SetUnbreakable(node.children[0]) |
| if not style.Get('DEDENT_CLOSING_BRACKETS'): |
| self._SetStronglyConnected(node.children[-1]) |
| |
| def Visit_argument(self, node): # pylint: disable=invalid-name |
| # argument ::= test [comp_for] | test '=' test # Really [keyword '='] test |
| self.DefaultNodeVisit(node) |
| |
| index = 0 |
| while index < len(node.children) - 1: |
| next_child = node.children[index + 1] |
| if isinstance(next_child, pytree.Leaf) and next_child.value == '=': |
| self._SetUnbreakable(_FirstChildNode(node.children[index + 1])) |
| self._SetUnbreakable(_FirstChildNode(node.children[index + 2])) |
| index += 1 |
| |
| def Visit_dotted_name(self, node): # pylint: disable=invalid-name |
| # dotted_name ::= NAME ('.' NAME)* |
| self._SetUnbreakableOnChildren(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) |
| if pytree_utils.NodeName(child) == 'COLON': |
| # This is a key to a dictionary. We don't want to split the key if at |
| # all possible. |
| self._SetStronglyConnected(child) |
| |
| 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', '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]) |
| if pytree_utils.NodeName(node.children[-1]) == 'RSQB': |
| # Don't split the ending bracket of a subscript list. |
| self._SetStronglyConnected(node.children[-1]) |
| |
| 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 Visit_subscript(self, node): # pylint: disable=invalid-name |
| # subscript ::= test | [test] ':' [test] [sliceop] |
| self._SetStronglyConnected(*node.children) |
| 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 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, |
| style.Get('SPLIT_PENALTY_BEFORE_IF_EXPR')) |
| self._SetStronglyConnected(*node.children[1:]) |
| self.DefaultNodeVisit(node) |
| |
| def Visit_not_test(self, node): # pylint: disable=invalid-name |
| # not_test ::= 'not' not_test | comparison |
| self.DefaultNodeVisit(node) |
| self._SetExpressionPenalty(node, NOT_TEST) |
| |
| def Visit_comparison(self, node): # pylint: disable=invalid-name |
| # comparison ::= expr (comp_op expr)* |
| self.DefaultNodeVisit(node) |
| self._SetExpressionPenalty(node, COMPARISON_EXPRESSION) |
| |
| def Visit_arith_expr(self, node): # pylint: disable=invalid-name |
| # arith_expr ::= term (('+'|'-') term)* |
| self.DefaultNodeVisit(node) |
| self._SetExpressionPenalty(node, ARITHMETIC_EXPRESSION) |
| |
| 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, |
| 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: |
| 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._SetStronglyConnected(*node.children[1:]) |
| |
| ############################################################################ |
| # Helper methods that set the annotations. |
| |
| def _SetUnbreakable(self, node): |
| """Set an UNBREAKABLE penalty annotation for the given node.""" |
| _RecAnnotate(node, pytree_utils.Annotation.SPLIT_PENALTY, UNBREAKABLE) |
| |
| def _SetStronglyConnected(self, *nodes): |
| """Set a STRONGLY_CONNECTED penalty annotation for the given nodes.""" |
| for node in nodes: |
| _RecAnnotate(node, pytree_utils.Annotation.SPLIT_PENALTY, |
| STRONGLY_CONNECTED) |
| |
| def _SetUnbreakableOnChildren(self, node): |
| """Set an UNBREAKABLE penalty annotation on children of node.""" |
| for child in node.children: |
| self.Visit(child) |
| start = 2 if hasattr(node.children[0], 'is_pseudo') else 1 |
| for i in py3compat.range(start, len(node.children)): |
| self._SetUnbreakable(node.children[i]) |
| |
| def _SetExpressionPenalty(self, node, penalty): |
| """Set an ARITHMETIC_EXPRESSION penalty annotation children nodes.""" |
| |
| 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) |
| |
| RecArithmeticExpression(node, _FirstChildNode(node)) |
| |
| |
| 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 _FirstChildNode(node): |
| if isinstance(node, pytree.Leaf): |
| return node |
| return _FirstChildNode(node.children[0]) |
| |
| |
| def _LastChildNode(node): |
| if isinstance(node, pytree.Leaf): |
| return node |
| return _LastChildNode(node.children[-1]) |
