From c26daa4fd50e2db13f38279ce27b51b0ae4479fe Mon Sep 17 00:00:00 2001
From: =?utf8?q?=C5=81ukasz=20Langa?= <lukasz@langa.pl>
Date: Thu, 15 Mar 2018 19:20:42 -0700
Subject: [PATCH] Don't split on for-loop variable unpacks

Fixes #23
---
 README.md           |  3 +++
 black.py            | 27 +++++++++++++++++++++++++++
 tests/expression.py |  5 +++++
 3 files changed, 35 insertions(+)

diff --git a/README.md b/README.md
index ebff26a..b3e6985 100644
--- a/README.md
+++ b/README.md
@@ -262,6 +262,9 @@ You can still try but prepare to be disappointed.
 
 * fixed invalid spacing of dots in relative imports (#6, #13)
 
+* fixed invalid splitting after comma on unpacked variables in for-loops
+  (#23)
+
 * fixed spurious space in parenthesized set expressions (#7)
 
 * fixed spurious space after opening parentheses and in default
diff --git a/black.py b/black.py
index e2c427c..774d91d 100644
--- a/black.py
+++ b/black.py
@@ -392,6 +392,8 @@ class Line:
     comments: Dict[LeafID, Leaf] = attrib(default=Factory(dict))
     bracket_tracker: BracketTracker = attrib(default=Factory(BracketTracker))
     inside_brackets: bool = attrib(default=False)
+    has_for: bool = attrib(default=False)
+    _for_loop_variable: bool = attrib(default=False, init=False)
 
     def append(self, leaf: Leaf, preformatted: bool = False) -> None:
         has_value = leaf.value.strip()
@@ -403,8 +405,10 @@ class Line:
             # imports, for which we only preserve newlines.
             leaf.prefix += whitespace(leaf)
         if self.inside_brackets or not preformatted:
+            self.maybe_decrement_after_for_loop_variable(leaf)
             self.bracket_tracker.mark(leaf)
             self.maybe_remove_trailing_comma(leaf)
+            self.maybe_increment_for_loop_variable(leaf)
             if self.maybe_adapt_standalone_comment(leaf):
                 return
 
@@ -508,6 +512,29 @@ class Line:
 
         return False
 
+    def maybe_increment_for_loop_variable(self, leaf: Leaf) -> bool:
+        """In a for loop, or comprehension, the variables are often unpacks.
+
+        To avoid splitting on the comma in this situation, we will increase
+        the depth of tokens between `for` and `in`.
+        """
+        if leaf.type == token.NAME and leaf.value == 'for':
+            self.has_for = True
+            self.bracket_tracker.depth += 1
+            self._for_loop_variable = True
+            return True
+
+        return False
+
+    def maybe_decrement_after_for_loop_variable(self, leaf: Leaf) -> bool:
+        # See `maybe_increment_for_loop_variable` above for explanation.
+        if self._for_loop_variable and leaf.type == token.NAME and leaf.value == 'in':
+            self.bracket_tracker.depth -= 1
+            self._for_loop_variable = False
+            return True
+
+        return False
+
     def maybe_adapt_standalone_comment(self, comment: Leaf) -> bool:
         """Hack a standalone comment to act as a trailing comment for line splitting.
 
diff --git a/tests/expression.py b/tests/expression.py
index 6870749..59e4211 100644
--- a/tests/expression.py
+++ b/tests/expression.py
@@ -63,6 +63,7 @@ str or None if (1 if True else 2) else str or bytes or None
 [((i ** 2) + j) for i in (1, 2, 3) for j in (1, 2, 3)]
 {i: 0 for i in (1, 2, 3)}
 {i: j for i, j in ((1, 'a'), (2, 'b'), (3, 'c'))}
+{k: v for k, v in this_is_a_very_long_variable_which_will_cause_a_trailing_comma_which_breaks_the_comprehension}
 Python3 > Python2 > COBOL
 Life is Life
 call()
@@ -188,6 +189,10 @@ str or None if (1 if True else 2) else str or bytes or None
 [((i ** 2) + j) for i in (1, 2, 3) for j in (1, 2, 3)]
 {i: 0 for i in (1, 2, 3)}
 {i: j for i, j in ((1, 'a'), (2, 'b'), (3, 'c'))}
+{
+    k: v
+    for k, v in this_is_a_very_long_variable_which_will_cause_a_trailing_comma_which_breaks_the_comprehension
+}
 Python3 > Python2 > COBOL
 Life is Life
 call()
-- 
2.39.5