Fix SPLIT_ALL_COMMA_SEPARATED_VALUES and SPLIT_ALL_TOP_LEVEL_COMMA_SE…

…PARATED_VALUES being too agressive for lambdas and unpacking. (#1173)

Co-authored-by: Bill Wendling <morbo@google.com>

CHANGELOG.md

@@ -9,6 +9,8 @@
 ### Fixed
 - Fix SPLIT_ARGUMENTS_WHEN_COMMA_TERMINATED for one-item named argument lists
   by taking precedence over SPLIT_BEFORE_NAMED_ASSIGNS.
+- Fix SPLIT_ALL_COMMA_SEPARATED_VALUES and SPLIT_ALL_TOP_LEVEL_COMMA_SEPARATED_VALUES
+  being too agressive for lambdas and unpacking.
 
 ## [0.40.2] 2023-09-22
 ### Changes

yapf/pytree/subtype_assigner.py

@@ -222,6 +222,11 @@ def Visit_power(self, node):  # pylint: disable=invalid-name
       if isinstance(child, pytree.Leaf) and child.value == '**':
         _AppendTokenSubtype(child, subtypes.BINARY_OPERATOR)
 
+  def Visit_lambdef(self, node):  # pylint: disable=invalid-name
+    # trailer: '(' [arglist] ')' | '[' subscriptlist ']' | '.' NAME
+    _AppendSubtypeRec(node, subtypes.LAMBDEF)
+    self.DefaultNodeVisit(node)
+
   def Visit_trailer(self, node):  # pylint: disable=invalid-name
     for child in node.children:
       self.Visit(child)

yapf/yapflib/format_decision_state.py

@@ -180,6 +180,10 @@ def MustSplit(self):
       return False
 
     if style.Get('SPLIT_ALL_COMMA_SEPARATED_VALUES') and previous.value == ',':
+      if (subtypes.COMP_FOR in current.subtypes or
+          subtypes.LAMBDEF in current.subtypes):
+        return False
+
       return True
 
     if (style.Get('FORCE_MULTILINE_DICT') and
@@ -188,6 +192,11 @@ def MustSplit(self):
 
     if (style.Get('SPLIT_ALL_TOP_LEVEL_COMMA_SEPARATED_VALUES') and
         previous.value == ','):
+
+      if (subtypes.COMP_FOR in current.subtypes or
+          subtypes.LAMBDEF in current.subtypes):
+        return False
+
       # Avoid breaking in a container that fits in the current line if possible
       opening = _GetOpeningBracket(current)
 

yapf/yapflib/subtypes.py

@@ -38,3 +38,4 @@
 SIMPLE_EXPRESSION = 22
 PARAMETER_START = 23
 PARAMETER_STOP = 24
+LAMBDEF = 25

yapftests/reformatter_basic_test.py

@@ -91,6 +91,30 @@ def foo(long_arg,
     """)
     llines = yapf_test_helper.ParseAndUnwrap(unformatted_code)
     self.assertCodeEqual(expected_formatted_code, reformatter.Reformat(llines))
+    unformatted_code = textwrap.dedent("""\
+        values = [ lambda arg1, arg2: arg1 + arg2 ]
+    """)  # noqa
+    expected_formatted_code = textwrap.dedent("""\
+        values = [
+            lambda arg1, arg2: arg1 + arg2
+        ]
+    """)
+    llines = yapf_test_helper.ParseAndUnwrap(unformatted_code)
+    self.assertCodeEqual(expected_formatted_code, reformatter.Reformat(llines))
+    unformatted_code = textwrap.dedent("""\
+        values = [
+            (some_arg1, some_arg2) for some_arg1, some_arg2 in values
+        ]
+    """)  # noqa
+    expected_formatted_code = textwrap.dedent("""\
+        values = [
+            (some_arg1,
+             some_arg2)
+            for some_arg1, some_arg2 in values
+        ]
+    """)
+    llines = yapf_test_helper.ParseAndUnwrap(unformatted_code)
+    self.assertCodeEqual(expected_formatted_code, reformatter.Reformat(llines))
     # There is a test for split_all_top_level_comma_separated_values, with
     # different expected value
     unformatted_code = textwrap.dedent("""\
@@ -161,6 +185,32 @@ def foo(long_arg,
     """)
     llines = yapf_test_helper.ParseAndUnwrap(unformatted_code)
     self.assertCodeEqual(expected_formatted_code, reformatter.Reformat(llines))
+    # Works the same way as split_all_comma_separated_values
+    unformatted_code = textwrap.dedent("""\
+        values = [ lambda arg1, arg2: arg1 + arg2 ]
+    """)  # noqa
+    expected_formatted_code = textwrap.dedent("""\
+        values = [
+            lambda arg1, arg2: arg1 + arg2
+        ]
+    """)
+    llines = yapf_test_helper.ParseAndUnwrap(unformatted_code)
+    self.assertCodeEqual(expected_formatted_code, reformatter.Reformat(llines))
+    # There is a test for split_all_comma_separated_values, with different
+    # expected value
+    unformatted_code = textwrap.dedent("""\
+        values = [
+            (some_arg1, some_arg2) for some_arg1, some_arg2 in values
+        ]
+    """)  # noqa
+    expected_formatted_code = textwrap.dedent("""\
+        values = [
+            (some_arg1, some_arg2)
+            for some_arg1, some_arg2 in values
+        ]
+    """)
+    llines = yapf_test_helper.ParseAndUnwrap(unformatted_code)
+    self.assertCodeEqual(expected_formatted_code, reformatter.Reformat(llines))
     # There is a test for split_all_comma_separated_values, with different
     # expected value
     unformatted_code = textwrap.dedent("""\