Adjust split penalties around arithmetic ops.

We want to break around the operator before the operands.

4 files changed, 38 insertions, 11 deletions

diff --git a/CHANGELOG b/CHANGELOG
index 699428b..904d01a 100644
--- a/CHANGELOG
+++ b/CHANGELOG
@@ -12,6 +12,8 @@
   without wrapping. If not, then we split.
 - Check all of the elements of a tuple. Similarly to how arguments are
   analyzed. This allows tuples to be split more rationally.
+- Adjust splitting penalties around arithmetic operators so that the code can
+  flow more freely. The code must flow!
 ### Fixed
 - Attempt to determine if long lambdas are allowed. This can be done on a
   case-by-case basis with a "pylint" disable comment.
diff --git a/yapf/yapflib/comment_splicer.py b/yapf/yapflib/comment_splicer.py
index 8717394..af999d2 100644
--- a/yapf/yapflib/comment_splicer.py
+++ b/yapf/yapflib/comment_splicer.py
@@ -177,8 +177,8 @@ def SpliceComments(tree):
                 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()))
+                    len(comment_prefix[rindex:]) - len(
+                        comment_prefix[rindex:].lstrip()))
                 comments = _CreateCommentsFromPrefix(
                     comment_prefix,
                     comment_lineno,
diff --git a/yapf/yapflib/split_penalty.py b/yapf/yapflib/split_penalty.py
index 66f733f..f815804 100644
--- a/yapf/yapflib/split_penalty.py
+++ b/yapf/yapflib/split_penalty.py
@@ -396,6 +396,8 @@ class _SplitPenaltyAssigner(pytree_visitor.PyTreeVisitor):
     self.DefaultNodeVisit(node)
     _IncreasePenalty(node, SHIFT_EXPR)
 
+  _ARITH_OPS = frozenset({'PLUS', 'MINUS'})
+
   def Visit_arith_expr(self, node):  # pylint: disable=invalid-name
     # arith_expr ::= term (('+'|'-') term)*
     self.DefaultNodeVisit(node)
@@ -404,19 +406,25 @@ class _SplitPenaltyAssigner(pytree_visitor.PyTreeVisitor):
     index = 1
     while index < len(node.children) - 1:
       child = node.children[index]
-      if isinstance(child, pytree.Leaf) and child.value in '+-':
+      if pytree_utils.NodeName(child) in self._ARITH_OPS:
         next_node = _FirstChildNode(node.children[index + 1])
-        _SetSplitPenalty(
-            next_node,
-            pytree_utils.GetNodeAnnotation(
-                next_node, pytree_utils.Annotation.SPLIT_PENALTY, default=0) -
-            100)
+        _SetSplitPenalty(next_node, ARITH_EXPR)
       index += 1
 
+  _TERM_OPS = frozenset({'STAR', 'AT', 'SLASH', 'PERCENT', 'DOUBLESLASH'})
+
   def Visit_term(self, node):  # pylint: disable=invalid-name
     # term ::= factor (('*'|'@'|'/'|'%'|'//') factor)*
-    _IncreasePenalty(node, TERM)
     self.DefaultNodeVisit(node)
+    _IncreasePenalty(node, TERM)
+
+    index = 1
+    while index < len(node.children) - 1:
+      child = node.children[index]
+      if pytree_utils.NodeName(child) in self._TERM_OPS:
+        next_node = _FirstChildNode(node.children[index + 1])
+        _SetSplitPenalty(next_node, TERM)
+      index += 1
 
   def Visit_factor(self, node):  # pyline: disable=invalid-name
     # factor ::= ('+'|'-'|'~') factor | power
diff --git a/yapftests/reformatter_buganizer_test.py b/yapftests/reformatter_buganizer_test.py
index 914befc..8f6ebac 100644
--- a/yapftests/reformatter_buganizer_test.py
+++ b/yapftests/reformatter_buganizer_test.py
@@ -28,6 +28,23 @@ class BuganizerFixes(yapf_test_helper.YAPFTest):
   def setUpClass(cls):
     style.SetGlobalStyle(style.CreateChromiumStyle())
 
+  def testB30500455(self):
+    unformatted_code = """\
+INITIAL_SYMTAB = dict([(name, 'exception#' + name) for name in INITIAL_EXCEPTIONS
+] * [(name, 'type#' + name) for name in INITIAL_TYPES] + [
+    (name, 'function#' + name) for name in INITIAL_FUNCTIONS
+] + [(name, 'const#' + name) for name in INITIAL_CONSTS])
+"""
+    expected_formatted_code = """\
+INITIAL_SYMTAB = dict(
+    [(name, 'exception#' + name) for name in INITIAL_EXCEPTIONS] *
+    [(name, 'type#' + name) for name in INITIAL_TYPES] +
+    [(name, 'function#' + name) for name in INITIAL_FUNCTIONS] +
+    [(name, 'const#' + name) for name in INITIAL_CONSTS])
+"""
+    uwlines = yapf_test_helper.ParseAndUnwrap(unformatted_code)
+    self.assertCodeEqual(expected_formatted_code, reformatter.Reformat(uwlines))
+
   def testB38343525(self):
     code = """\
 # This does foo.
@@ -503,8 +520,8 @@ def _():
     expected_formatted_code = textwrap.dedent("""\
         class _():
           def _():
-            hints.append(('hg tag -f -l -r %s %s # %s' %
-                          (short(ctx.node()), candidatetag, firstline))[:78])
+            hints.append(('hg tag -f -l -r %s %s # %s' % (short(
+                ctx.node()), candidatetag, firstline))[:78])
         """)
     uwlines = yapf_test_helper.ParseAndUnwrap(unformatted_code)
     self.assertCodeEqual(expected_formatted_code, reformatter.Reformat(uwlines))