Don't split on for-loop variable unpacks

Fixes #23
This commit is contained in:
Łukasz Langa 2018-03-15 19:20:42 -07:00 committed by Lukasz Langa
parent c85cb48ec2
commit c26daa4fd5
3 changed files with 35 additions and 0 deletions

View File

@ -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

View File

@ -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 @@ def append(self, leaf: Leaf, preformatted: bool = False) -> None:
# 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 @@ def maybe_remove_trailing_comma(self, closing: Leaf) -> bool:
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.

View File

@ -63,6 +63,7 @@
[((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 @@ async def f():
[((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()