Adjust bitwise operand split penalties for & and ^

Closes #665
google · Feb 10, 2019 · e2ce893 · e2ce893
1 parent a4d0b6d
commit e2ce893
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 27 deletions.
diff --git a/CHANGELOG b/CHANGELOG
@@ -6,6 +6,7 @@
 ### Fixed
 - Parse integer lists correctly, removing quotes if the list is within a
   string.
+- Adjust the penalties of bitwise operands for '&' and '^', similar to '|'.
 
 ## [0.26.0] 2019-02-08
 ### Added

diff --git a/yapf/yapflib/split_penalty.py b/yapf/yapflib/split_penalty.py
@@ -414,27 +414,19 @@ def Visit_expr(self, node):  # pylint: disable=invalid-name
     # expr ::= xor_expr ('|' xor_expr)*
     self.DefaultNodeVisit(node)
     _IncreasePenalty(node, EXPR)
-    index = 1
-    while index < len(node.children) - 1:
-      child = node.children[index]
-      if isinstance(child, pytree.Leaf) and child.value == '|':
-        if style.Get('SPLIT_BEFORE_BITWISE_OPERATOR'):
-          _SetSplitPenalty(child, style.Get('SPLIT_PENALTY_BITWISE_OPERATOR'))
-        else:
-          _SetSplitPenalty(
-              pytree_utils.FirstLeafNode(node.children[index + 1]),
-              style.Get('SPLIT_PENALTY_BITWISE_OPERATOR'))
-      index += 1
+    _SetBitwiseOperandPenalty(node, '|')
 
   def Visit_xor_expr(self, node):  # pylint: disable=invalid-name
     # xor_expr ::= and_expr ('^' and_expr)*
     self.DefaultNodeVisit(node)
     _IncreasePenalty(node, XOR_EXPR)
+    _SetBitwiseOperandPenalty(node, '^')
 
   def Visit_and_expr(self, node):  # pylint: disable=invalid-name
     # and_expr ::= shift_expr ('&' shift_expr)*
     self.DefaultNodeVisit(node)
     _IncreasePenalty(node, AND_EXPR)
+    _SetBitwiseOperandPenalty(node, '&')
 
   def Visit_shift_expr(self, node):  # pylint: disable=invalid-name
     # shift_expr ::= arith_expr (('<<'|'>>') arith_expr)*
@@ -447,29 +439,15 @@ def Visit_arith_expr(self, node):  # pylint: disable=invalid-name
     # arith_expr ::= term (('+'|'-') term)*
     self.DefaultNodeVisit(node)
     _IncreasePenalty(node, ARITH_EXPR)
-
-    index = 1
-    while index < len(node.children) - 1:
-      child = node.children[index]
-      if pytree_utils.NodeName(child) in self._ARITH_OPS:
-        next_node = pytree_utils.FirstLeafNode(node.children[index + 1])
-        _SetSplitPenalty(next_node, ARITH_EXPR)
-      index += 1
+    _SetExpressionOperandPenalty(node, self._ARITH_OPS, ARITH_EXPR)
 
   _TERM_OPS = frozenset({'STAR', 'AT', 'SLASH', 'PERCENT', 'DOUBLESLASH'})
 
   def Visit_term(self, node):  # pylint: disable=invalid-name
     # term ::= factor (('*'|'@'|'/'|'%'|'//') factor)*
     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 = pytree_utils.FirstLeafNode(node.children[index + 1])
-        _SetSplitPenalty(next_node, TERM)
-      index += 1
+    _SetExpressionOperandPenalty(node, self._TERM_OPS, TERM)
 
   def Visit_factor(self, node):  # pyline: disable=invalid-name
     # factor ::= ('+'|'-'|'~') factor | power
@@ -558,6 +536,26 @@ def RecExpression(node, first_child_leaf):
   RecExpression(node, pytree_utils.FirstLeafNode(node))
 
 
+def _SetBitwiseOperandPenalty(node, op):
+  for index in py3compat.range(1, len(node.children) - 1):
+    child = node.children[index]
+    if isinstance(child, pytree.Leaf) and child.value == op:
+      if style.Get('SPLIT_BEFORE_BITWISE_OPERATOR'):
+        _SetSplitPenalty(child, style.Get('SPLIT_PENALTY_BITWISE_OPERATOR'))
+      else:
+        _SetSplitPenalty(
+            pytree_utils.FirstLeafNode(node.children[index + 1]),
+            style.Get('SPLIT_PENALTY_BITWISE_OPERATOR'))
+
+
+def _SetExpressionOperandPenalty(node, ops, penalty):
+  for index in py3compat.range(1, len(node.children) - 1):
+      child = node.children[index]
+      if pytree_utils.NodeName(child) in ops:
+        next_node = pytree_utils.FirstLeafNode(node.children[index + 1])
+        _SetSplitPenalty(next_node, penalty)
+
+
 def _IncreasePenalty(node, amt):
   """Increase a penalty annotation on children nodes."""
 

diff --git a/yapftests/reformatter_pep8_test.py b/yapftests/reformatter_pep8_test.py
@@ -431,6 +431,29 @@ def testNoSplitBeforeDictValue(self):
     finally:
       style.SetGlobalStyle(style.CreatePEP8Style())
 
+  def testBitwiseOperandSplitting(self):
+    unformatted_code = """\
+def _():
+    include_values = np.where(
+                (cdffile['Quality_Flag'][:] >= 5) & (
+                cdffile['Day_Night_Flag'][:] == 1) & (
+                cdffile['Longitude'][:] >= select_lon - radius) & (
+                cdffile['Longitude'][:] <= select_lon + radius) & (
+                cdffile['Latitude'][:] >= select_lat - radius) & (
+                cdffile['Latitude'][:] <= select_lat + radius))
+"""
+    expected_code = """\
+def _():
+    include_values = np.where(
+        (cdffile['Quality_Flag'][:] >= 5) & (cdffile['Day_Night_Flag'][:] == 1)
+        & (cdffile['Longitude'][:] >= select_lon - radius)
+        & (cdffile['Longitude'][:] <= select_lon + radius)
+        & (cdffile['Latitude'][:] >= select_lat - radius)
+        & (cdffile['Latitude'][:] <= select_lat + radius))
+"""
+    uwlines = yapf_test_helper.ParseAndUnwrap(unformatted_code)
+    self.assertEqual(expected_code, reformatter.Reformat(uwlines))
+
 
 if __name__ == '__main__':
   unittest.main()