Repository: google/yapf
Pull request #1062 (2023-02-24 20:29:18)

Generalized the ending comma heuristic to subscripts
This commit ensures that

    G[A, B, C,]

is split in a similar way to

    [A, B, C,]

i.e.

    G[
        A,
        B,
        C,
     ]

instead of

    G[A, B,
      C,]

assuming A, B and C do not fit on a single line.

Long subscript expressions often occur when instantiating generic classes, and having YAPF reformat them to a more readable multi-line variant seems like a "good thing".

---

Note that the heuristic is not perfect and in particular could produce unexpected/verbose outputs when one of the items is itself a list with an ending comma, e.g.

    G[[A,], B, C,]

is reformatted to

    G[
        [
            A,
        ],
        B,
        C,
    ]

where G could be typing.Callable.


Patch:
diff --git a/CHANGELOG b/CHANGELOG
index f3ba6e9..467c491 100644
--- a/CHANGELOG
+++ b/CHANGELOG
@@ -10,6 +10,7 @@
 - Moved 'pytree' parsing tools into its own subdirectory.
 - Add support for Python 3.10.
 - Format generated dicts with respect to same rules as regular dicts
+- Generalized the ending comma heuristic to subscripts.
 ### Fixed
 - Split line before all comparison operators.
 
diff --git a/yapf/pytree/pytree_unwrapper.py b/yapf/pytree/pytree_unwrapper.py
index 3fe4ade..ba1e0c4 100644
--- a/yapf/pytree/pytree_unwrapper.py
+++ b/yapf/pytree/pytree_unwrapper.py
@@ -285,6 +285,10 @@ class PyTreeUnwrapper(pytree_visitor.PyTreeVisitor):
     _DetermineMustSplitAnnotation(node)
     self.DefaultNodeVisit(node)
 
+  def Visit_subscriptlist(self, node):  # pylint: disable=invalid-name
+    _DetermineMustSplitAnnotation(node)
+    self.DefaultNodeVisit(node)
+
   def DefaultLeafVisit(self, leaf):
     """Default visitor for tree leaves.
 
diff --git a/yapf/yapflib/format_decision_state.py b/yapf/yapflib/format_decision_state.py
index edab8ad..2468b1f 100644
--- a/yapf/yapflib/format_decision_state.py
+++ b/yapf/yapflib/format_decision_state.py
@@ -205,7 +205,9 @@ class FormatDecisionState(object):
         (current.value in '}]' and style.Get('SPLIT_BEFORE_CLOSING_BRACKET') or
          current.value in '}])' and style.Get('INDENT_CLOSING_BRACKETS'))):
       # Split before the closing bracket if we can.
-      if subtypes.SUBSCRIPT_BRACKET not in current.subtypes:
+      if (subtypes.SUBSCRIPT_BRACKET not in current.subtypes or
+          (previous.value == ',' and
+           not style.Get('DISABLE_ENDING_COMMA_HEURISTIC'))):
         return current.node_split_penalty != split_penalty.UNBREAKABLE
 
     if (current.value == ')' and previous.value == ',' and


Test patch:
diff --git a/yapftests/reformatter_basic_test.py b/yapftests/reformatter_basic_test.py
index 935c7c3..e07d46d 100644
--- a/yapftests/reformatter_basic_test.py
+++ b/yapftests/reformatter_basic_test.py
@@ -2645,6 +2645,19 @@ class A(object):
     llines = yapf_test_helper.ParseAndUnwrap(code)
     self.assertCodeEqual(code, reformatter.Reformat(llines))
 
+  def testSubscriptExpressionTerminatedByComma(self):
+    unformatted_code = textwrap.dedent("""\
+        A[B, C,]
+        """)
+    expected_code = textwrap.dedent("""\
+        A[
+            B,
+            C,
+        ]
+    """)
+    llines = yapf_test_helper.ParseAndUnwrap(unformatted_code)
+    self.assertCodeEqual(expected_code, reformatter.Reformat(llines))
+
   def testListWithFunctionCalls(self):
     unformatted_code = textwrap.dedent("""\
         def foo():