split_penalty.py

# 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])
    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
    self._SetUnbreakable(node.children[colon_idx])
    self.DefaultNodeVisit(node)
    if arrow_idx > 0:
      pytree_utils.SetNodeAnnotation(
          _LastChildNode(node.children[arrow_idx - 1]),
          pytree_utils.Annotation.SPLIT_PENALTY, 0)
      self._SetUnbreakable(node.children[arrow_idx])
      self._SetStronglyConnected(node.children[arrow_idx + 1])

  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 a penalty annotation on 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])