Add knob disabling comma ending heuristic

5 files changed, 31 insertions, 1 deletions

diff --git a/CHANGELOG b/CHANGELOG
index 767043a..4693f03 100644
--- a/CHANGELOG
+++ b/CHANGELOG
@@ -2,7 +2,10 @@
 # All notable changes to this project will be documented in this file.
 # This project adheres to [Semantic Versioning](http://semver.org/).
 
-## [0.22.1] UNRELEASED
+## [0.23.0] UNRELEASED
+### Added
+- `DISABLE_ENDING_COMMA_HEURISTIC` is a new knob to disable the heuristic which
+  splits a list onto separate lines if the list is comma-terminated.
 ### Fixed
 - There's no need to increase N_TOKENS. In fact, it causes other things which
   use lib2to3 to fail if called from YAPF.
diff --git a/README.rst b/README.rst
index 120acac..19ef0c2 100644
--- a/README.rst
+++ b/README.rst
@@ -415,6 +415,10 @@ Knobs
             end_ts=now(),
         )  # <--- this bracket is dedented and on a separate line
 
+``DISABLE_ENDING_COMMA_HEURISTIC``
+    Disable the heuristic which places each list element on a separate line if
+    the list is comma-terminated.
+
 ``EACH_DICT_ENTRY_ON_SEPARATE_LINE``
     Place each dictionary entry onto its own line.
 
diff --git a/yapf/yapflib/pytree_unwrapper.py b/yapf/yapflib/pytree_unwrapper.py
index 550bc89..0d371ae 100644
--- a/yapf/yapflib/pytree_unwrapper.py
+++ b/yapf/yapflib/pytree_unwrapper.py
@@ -34,6 +34,7 @@ from lib2to3.pgen2 import token as grammar_token
 from yapf.yapflib import pytree_utils
 from yapf.yapflib import pytree_visitor
 from yapf.yapflib import split_penalty
+from yapf.yapflib import style
 from yapf.yapflib import unwrapped_line
 
 
@@ -342,6 +343,8 @@ def _AdjustSplitPenalty(uwline):
 
 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):
     token = next(node.parent.leaves())
     if token.value == '(':
diff --git a/yapf/yapflib/style.py b/yapf/yapflib/style.py
index c0920c5..6144246 100644
--- a/yapf/yapflib/style.py
+++ b/yapf/yapflib/style.py
@@ -129,6 +129,9 @@ _STYLE_HELP = dict(
             start_ts=now()-timedelta(days=3),
             end_ts=now(),
         )        # <--- this bracket is dedented and on a separate line"""),
+    DISABLE_ENDING_COMMA_HEURISTIC=textwrap.dedent("""\
+      Disable the heuristic which places each list element on a separate line
+      if the list is comma-terminated."""),
     EACH_DICT_ENTRY_ON_SEPARATE_LINE=textwrap.dedent("""\
       Place each dictionary entry onto its own line."""),
     I18N_COMMENT=textwrap.dedent("""\
@@ -272,6 +275,7 @@ def CreatePEP8Style():
       CONTINUATION_ALIGN_STYLE='SPACE',
       CONTINUATION_INDENT_WIDTH=4,
       DEDENT_CLOSING_BRACKETS=False,
+      DISABLE_ENDING_COMMA_HEURISTIC=False,
       EACH_DICT_ENTRY_ON_SEPARATE_LINE=True,
       I18N_COMMENT='',
       I18N_FUNCTION_CALL='',
@@ -420,6 +424,7 @@ _STYLE_OPTION_VALUE_CONVERTER = dict(
     CONTINUATION_ALIGN_STYLE=_ContinuationAlignStyleStringConverter,
     CONTINUATION_INDENT_WIDTH=int,
     DEDENT_CLOSING_BRACKETS=_BoolConverter,
+    DISABLE_ENDING_COMMA_HEURISTIC=_BoolConverter,
     EACH_DICT_ENTRY_ON_SEPARATE_LINE=_BoolConverter,
     I18N_COMMENT=str,
     I18N_FUNCTION_CALL=_StringListConverter,
diff --git a/yapftests/reformatter_basic_test.py b/yapftests/reformatter_basic_test.py
index 360aae4..2078b09 100644
--- a/yapftests/reformatter_basic_test.py
+++ b/yapftests/reformatter_basic_test.py
@@ -2490,6 +2490,21 @@ s = 'foo \\
     finally:
       style.SetGlobalStyle(style.CreateChromiumStyle())
 
+  def testDisableEndingCommaHeuristic(self):
+    try:
+      style.SetGlobalStyle(
+          style.CreateStyleFromConfig(
+              '{based_on_style: chromium,'
+              ' disable_ending_comma_heuristic: True}'))
+
+      code = """\
+x = [1, 2, 3, 4, 5, 6, 7,]
+"""
+      uwlines = yapf_test_helper.ParseAndUnwrap(code)
+      self.assertCodeEqual(code, reformatter.Reformat(uwlines))
+    finally:
+      style.SetGlobalStyle(style.CreateChromiumStyle())
+
 
 if __name__ == '__main__':
   unittest.main()