From b7991b310541294f18f90e52f6fcd4d7c35541b6 Mon Sep 17 00:00:00 2001
From: Jesse Clemens <jesse.clemens@sony.com>
Date: Wed, 18 Oct 2023 14:40:23 -0400
Subject: [PATCH] Allow containers to wrap gracefully

And do not force breaks on commas in the middle of for comprehensions
---
 yapf/yapflib/format_decision_state.py | 19 ++++++++++++-------
 1 file changed, 12 insertions(+), 7 deletions(-)

diff --git a/yapf/yapflib/format_decision_state.py b/yapf/yapflib/format_decision_state.py
index 4817c924b..441915973 100644
--- a/yapf/yapflib/format_decision_state.py
+++ b/yapf/yapflib/format_decision_state.py
@@ -203,7 +203,9 @@ def MustSplit(self):
       # Allow the fallthrough code to handle the closing bracket.
       if current != opening.matching_bracket:
         # If the container doesn't fit in the current line, must split
-        return not self._ContainerFitsOnStartLine(opening)
+        if (subtypes.COMP_FOR not in current.subtypes and
+            not self._ContainerFitsOnStartLine(opening)):
+          return True
 
     if (self.stack[-1].split_before_closing_bracket and
         (current.value in '}]' and style.Get('SPLIT_BEFORE_CLOSING_BRACKET') or
@@ -546,9 +548,14 @@ def SurroundedByParens(token):
             return True
       else:
         # Split after the opening of a container if it doesn't fit on the
-        # current line.
+        # current line, including checking for wrapping.
         if not self._FitsOnLine(previous, previous.matching_bracket):
-          return True
+          arg_lengths = _CalculateArgLengths(previous)
+          start_col = self.column + len(current.value) + len(previous.value)
+          for length in arg_lengths:
+            length += start_col
+            if length > self.column_limit:
+              return True
 
     ###########################################################################
     # Original Formatting Splitting
@@ -1106,10 +1113,8 @@ def _ContainerFitsOnStartLine(self, opening):
             self.stack[-1].indent) <= self.column_limit
 
 
-_COMPOUND_STMTS = frozenset({
-    'for', 'while', 'if', 'elif', 'with', 'except', 'def', 'class', 'match',
-    'case'
-})
+_COMPOUND_STMTS = frozenset({'for', 'while', 'if', 'elif', 'with', 'except',
+                             'def', 'class', 'match', 'case'})
 
 
 def _IsCompoundStatement(token):