Add new DISABLE_SPLIT_LIST_WITH_COMMENT flag.

CHANGELOG.md

@@ -6,6 +6,28 @@
 ### Changes
 - Remove dependency on importlib-metadata
 - Remove dependency on tomli when using >= py311
+### Added
+- `DISABLE_SPLIT_LIST_WITH_COMMENT` is a new knob that changes the
+  behavior of splitting a list when a comment is present inside the list.
+
+  Before, we split a list containing a comment just like we split a list
+  containing a trailing comma: Each element goes on its own line (unless
+  `DISABLE_ENDING_COMMA_HEURISTIC` is true).
+
+  Now, if `DISABLE_SPLIT_LIST_WITH_COMMENT` is true, we do not split every
+  element of the list onto a new line just because there's a comment somewhere
+  in the list.
+
+  This mirrors the behavior of clang-format, and is useful for e.g. forming
+  "logical groups" of elements in a list.
+
+  Note: Upgrading will result in a behavioral change if you have
+  `DISABLE_ENDING_COMMA_HEURISTIC` in your config.  Before this version, this
+  flag caused us not to split lists with a trailing comma *and* lists that
+  contain comments.  Now, if you set only that flag, we *will* split lists
+  that contain comments.  Set the new `DISABLE_SPLIT_LIST_WITH_COMMENT` flag to
+  true to preserve the old behavior.
+
 ### Fixed
 - Fix SPLIT_ARGUMENTS_WHEN_COMMA_TERMINATED for one-item named argument lists
   by taking precedence over SPLIT_BEFORE_NAMED_ASSIGNS.

README.md

@@ -513,6 +513,33 @@ optional arguments:
 >    Disable the heuristic which places each list element on a separate line if
 >    the list is comma-terminated.
 
+#### `DISABLE_DISABLE_SPLIT_LIST_WITH_COMMENT`
+
+>    Don't put every element on a new line within a list that contains
+>    interstitial comments.
+>
+>    Without this flag (default):
+>
+>    ```
+>    [
+>      a,
+>      b,  #
+>      c
+>    ]
+>    ```
+>
+>    With this flag:
+>
+>    ```
+>    [
+>      a, b,  #
+>      c
+>    ]
+>    ```
+>
+>    This is useful for forming "logical groups" of elements in a list.  It also
+>    works in function declarations.
+
 #### `EACH_DICT_ENTRY_ON_SEPARATE_LINE`
 
 >    Place each dictionary entry onto its own line.

yapf/pytree/pytree_unwrapper.py

@@ -407,16 +407,27 @@ def _AdjustSplitPenalty(line):
 
 def _DetermineMustSplitAnnotation(node):
   """Enforce a split in the list if the list ends with a comma."""
-  if style.Get('DISABLE_ENDING_COMMA_HEURISTIC'):
-    return
-  if not _ContainsComments(node):
+
+  def SplitBecauseTrailingComma():
+    if style.Get('DISABLE_ENDING_COMMA_HEURISTIC'):
+      return False
     token = next(node.parent.leaves())
     if token.value == '(':
       if sum(1 for ch in node.children if ch.type == grammar_token.COMMA) < 2:
-        return
+        return False
     if (not isinstance(node.children[-1], pytree.Leaf) or
         node.children[-1].value != ','):
-      return
+      return False
+    return True
+
+  def SplitBecauseListContainsComment():
+    return (not style.Get('DISABLE_SPLIT_LIST_WITH_COMMENT') and
+            _ContainsComments(node))
+
+  if (not SplitBecauseTrailingComma() and
+      not SplitBecauseListContainsComment()):
+    return
+
   num_children = len(node.children)
   index = 0
   _SetMustSplitOnFirstLeaf(node.children[0])

yapf/yapflib/style.py

@@ -180,6 +180,10 @@ def method():
       Disable the heuristic which places each list element on a separate line
       if the list is comma-terminated.
     """),
+    DISABLE_SPLIT_LIST_WITH_COMMENT=textwrap.dedent("""
+      Don't put every element on a new line within a list that contains
+      interstitial comments.
+    """),
     EACH_DICT_ENTRY_ON_SEPARATE_LINE=textwrap.dedent("""\
       Place each dictionary entry onto its own line.
     """),
@@ -483,6 +487,7 @@ def CreatePEP8Style():
       CONTINUATION_INDENT_WIDTH=4,
       DEDENT_CLOSING_BRACKETS=False,
       DISABLE_ENDING_COMMA_HEURISTIC=False,
+      DISABLE_SPLIT_LIST_WITH_COMMENT=False,
       EACH_DICT_ENTRY_ON_SEPARATE_LINE=True,
       FORCE_MULTILINE_DICT=False,
       I18N_COMMENT='',
@@ -671,6 +676,7 @@ def _IntOrIntListConverter(s):
     CONTINUATION_INDENT_WIDTH=int,
     DEDENT_CLOSING_BRACKETS=_BoolConverter,
     DISABLE_ENDING_COMMA_HEURISTIC=_BoolConverter,
+    DISABLE_SPLIT_LIST_WITH_COMMENT=_BoolConverter,
     EACH_DICT_ENTRY_ON_SEPARATE_LINE=_BoolConverter,
     FORCE_MULTILINE_DICT=_BoolConverter,
     I18N_COMMENT=str,

yapftests/reformatter_basic_test.py

@@ -345,6 +345,29 @@ def f(  # Intermediate comment
     llines = yapf_test_helper.ParseAndUnwrap(unformatted_code)
     self.assertCodeEqual(expected_formatted_code, reformatter.Reformat(llines))
 
+  def testParamListWithTrailingComments(self):
+    unformatted_code = textwrap.dedent("""\
+        def f(a,
+              b, #
+              c):
+          pass
+    """)
+    expected_formatted_code = textwrap.dedent("""\
+        def f(a, b,  #
+              c):
+          pass
+    """)
+    try:
+      style.SetGlobalStyle(
+          style.CreateStyleFromConfig(
+              '{based_on_style: yapf,'
+              ' disable_split_list_with_comment: True}'))
+      llines = yapf_test_helper.ParseAndUnwrap(unformatted_code)
+      self.assertCodeEqual(expected_formatted_code,
+                           reformatter.Reformat(llines))
+    finally:
+      style.SetGlobalStyle(style.CreateYapfStyle())
+
   def testBlankLinesBetweenTopLevelImportsAndVariables(self):
     unformatted_code = textwrap.dedent("""\
         import foo as bar