Repository: google/yapf
Pull request #470 (2017-10-29 23:42:50)

Improve formatting for comprehensions and generators.
Terminology reference:
```
[loop_var for loop_var in a_list if loop_var]
 ^      ^ ^                    ^ ^         ^
 \-expr-/ \------comp_for------/ \-comp_if-/
```

Implements the following preferences:
- Try to keep comprehensions on a single line.
- Prefer splitting at comp_if when comp_for was split as well.
- Try to avoid splitting at comp_for when expr is exactly one token long
  - This keeps the (often short) clauses together: [a for a in ...].

Also: Fixes a bug in the penalty logic where comprehensions provided as
arguments to functions had a heavy penalty for splitting at their start.

Previous Behavior:

```python
obj.FunctionCall(loop_var.Get()
                 for loop_var in a_very_long_list_name)
```

Fixed Example:

```python
obj.FunctionCall(
    loop_var.Get() for loop_var in a_very_long_list_name)
```



Patch:
diff --git a/CHANGELOG b/CHANGELOG
index 070c45d..fa127c8 100644
--- a/CHANGELOG
+++ b/CHANGELOG
@@ -3,6 +3,11 @@
 # This project adheres to [Semantic Versioning](http://semver.org/).
 
 ## [0.19.1] UNRELEASED
+### Added
+- Improve splitting of comprehensions and generators. Add
+  `SPLIT_PENALTY_COMPREHENSION` knob to control preference for keeping
+  comprehensions on a single line and `SPLIT_COMPLEX_COMPREHENSION` to enable
+  splitting each clause of complex comprehensions onto its own line.
 ### Changed
 - Take into account a named function argument when determining if we should
   split before the first argument in a function call.
@@ -19,6 +24,8 @@
   doesn't it on the current line. We try hard not to move such comments around.
 - Use a TextIOWrapper when reading from stdin in Python3. This is necessary for
   some encodings, like cp936, used on Windows.
+- Remove the penalty for a split before the first argument in a function call
+  where the only argument is a generator expression.
 
 ## [0.19.0] 2017-10-14
 ### Added
diff --git a/README.rst b/README.rst
index ae2bc74..96a648d 100644
--- a/README.rst
+++ b/README.rst
@@ -461,6 +461,27 @@ Knobs
 ``SPLIT_BEFORE_NAMED_ASSIGNS``
     Split named assignments onto individual lines.
 
+``SPLIT_COMPLEX_COMPREHENSION``
+    For list comprehensions and generator expressions with multiple clauses
+    (e.g mutiple "for" calls, "if" filter expressions) and which need to be
+    reflowed, split each clause onto its own line. For example:
+
+    .. code-block:: python
+
+      result = [
+          a_var + b_var for a_var in xrange(1000) for b_var in xrange(1000)
+          if a_var % b_var]
+
+    would reformat to something like:
+
+    .. code-block:: python
+
+      result = [
+          a_var + b_var
+          for a_var in xrange(1000)
+          for b_var in xrange(1000)
+          if a_var % b_var]
+
 ``SPLIT_PENALTY_AFTER_OPENING_BRACKET``
     The penalty for splitting right after the opening bracket.
 
@@ -474,6 +495,9 @@ Knobs
     The penalty of splitting the line around the ``&``, ``|``, and ``^``
     operators.
 
+``SPLIT_PENALTY_COMPREHENSION``
+    The penalty for splitting a list comprehension or generator expression.
+
 ``SPLIT_PENALTY_EXCESS_CHARACTER``
     The penalty for characters over the column limit.
 
diff --git a/yapf/yapflib/file_resources.py b/yapf/yapflib/file_resources.py
index 021764e..063161b 100644
--- a/yapf/yapflib/file_resources.py
+++ b/yapf/yapflib/file_resources.py
@@ -121,7 +121,8 @@ def _FindPythonFiles(filenames, recursive, exclude):
         # TODO(morbo): Look into a version of os.walk that can handle recursion.
         python_files.extend(
             os.path.join(dirpath, f)
-            for dirpath, _, filelist in os.walk(filename) for f in filelist
+            for dirpath, _, filelist in os.walk(filename)
+            for f in filelist
             if IsPythonFile(os.path.join(dirpath, f)))
       else:
         raise errors.YapfError(
diff --git a/yapf/yapflib/format_decision_state.py b/yapf/yapflib/format_decision_state.py
index de781c2..8d18c54 100644
--- a/yapf/yapflib/format_decision_state.py
+++ b/yapf/yapflib/format_decision_state.py
@@ -49,6 +49,8 @@ class FormatDecisionState(object):
     previous: The previous format decision state in the decision tree.
     stack: A stack (of _ParenState) keeping track of properties applying to
       parenthesis levels.
+    comp_stack: A stack (of _ComprehensionState) keeping track of properties
+      applying to comprehensions.
     ignore_stack_for_comparison: Ignore the stack of _ParenState for state
       comparison.
   """
@@ -71,6 +73,7 @@ class FormatDecisionState(object):
     self.lowest_level_on_line = 0
     self.ignore_stack_for_comparison = False
     self.stack = [_ParenState(first_indent, first_indent)]
+    self.comp_stack = []
     self.first_indent = first_indent
     self.newline = False
     self.previous = None
@@ -90,6 +93,7 @@ class FormatDecisionState(object):
     new.newline = self.newline
     new.previous = self.previous
     new.stack = [state.Clone() for state in self.stack]
+    new.comp_stack = [state.Clone() for state in self.comp_stack]
     return new
 
   def __eq__(self, other):
@@ -102,7 +106,8 @@ class FormatDecisionState(object):
             self.start_of_line_level == other.start_of_line_level and
             self.lowest_level_on_line == other.lowest_level_on_line and
             (self.ignore_stack_for_comparison or
-             other.ignore_stack_for_comparison or self.stack == other.stack))
+             other.ignore_stack_for_comparison or
+             self.stack == other.stack and self.comp_stack == other.comp_stack))
 
   def __ne__(self, other):
     return not self == other
@@ -498,6 +503,8 @@ class FormatDecisionState(object):
     else:
       self._AddTokenOnCurrentLine(dry_run)
 
+    penalty += self._CalculateComprehensionState(newline)
+
     return self.MoveStateToNextToken() + penalty
 
   def _AddTokenOnCurrentLine(self, dry_run):
@@ -637,6 +644,76 @@ class FormatDecisionState(object):
 
     return top_of_stack.indent
 
+  def _CalculateComprehensionState(self, newline):
+    """Makes required changes to comprehension state.
+
+    Args:
+      newline: Whether the current token is to be added on a newline.
+
+    Returns:
+      The penalty for the token-newline combination given the current
+      comprehension state.
+    """
+    current = self.next_token
+    previous = current.previous_token
+    top_of_stack = self.comp_stack[-1] if self.comp_stack else None
+    penalty = 0
+
+    if top_of_stack is not None:
+      # Check if the token terminates the current comprehension.
+      if current == top_of_stack.closing_bracket:
+        last = self.comp_stack.pop()
+        # Lightly penalize comprehensions that are split across multiple lines.
+        if last.has_interior_split:
+          penalty += style.Get('SPLIT_PENALTY_COMPREHENSION')
+
+        return penalty
+
+      if newline:
+        top_of_stack.has_interior_split = True
+
+    if (format_token.Subtype.COMP_EXPR in current.subtypes and
+        format_token.Subtype.COMP_EXPR not in previous.subtypes):
+      self.comp_stack.append(_ComprehensionState(current))
+
+      return penalty
+
+    if (current.value == 'for' and
+        format_token.Subtype.COMP_FOR in current.subtypes):
+      if top_of_stack.for_token is not None:
+        # Treat nested comprehensions like normal comp_if expressions.
+        # Example:
+        #     my_comp = [
+        #         a.qux + b.qux
+        #         for a in foo
+        #   -->   for b in bar   <--
+        #         if a.zut + b.zut
+        #     ]
+        if (style.Get('SPLIT_COMPLEX_COMPREHENSION') and
+            top_of_stack.has_split_at_for != newline and
+            (top_of_stack.has_split_at_for or
+             not top_of_stack.HasTrivialExpr())):
+          penalty += split_penalty.UNBREAKABLE
+      else:
+        top_of_stack.for_token = current
+        top_of_stack.has_split_at_for = newline
+
+        # Try to keep trivial expressions on the same line as the comp_for.
+        if (style.Get('SPLIT_COMPLEX_COMPREHENSION') and newline and
+            top_of_stack.HasTrivialExpr()):
+          penalty += split_penalty.CONNECTED
+
+    if (format_token.Subtype.COMP_IF in current.subtypes and
+        format_token.Subtype.COMP_IF not in previous.subtypes):
+      # Penalize breaking at comp_if when it doesn't match the newline structure
+      # in the rest of the comprehension.
+      if (style.Get('SPLIT_COMPLEX_COMPREHENSION') and
+          top_of_stack.has_split_at_for != newline and
+          (top_of_stack.has_split_at_for or not top_of_stack.HasTrivialExpr())):
+        penalty += split_penalty.UNBREAKABLE
+
+    return penalty
+
   def MoveStateToNextToken(self):
     """Calculate format decision state information and move onto the next token.
 
@@ -915,3 +992,60 @@ class _ParenState(object):
   def __hash__(self, *args, **kwargs):
     return hash((self.indent, self.last_space, self.closing_scope_indent,
                  self.split_before_closing_bracket, self.num_line_splits))
+
+
+class _ComprehensionState(object):
+  """Maintains the state of list comprehension line-break decisions.
+
+  A stack of _ComprehensionState objects are kept to ensure that list
+  comprehensions are wrapped with well-defined rules.
+
+  Attributes:
+    expr_token: The first token in the comprehension.
+    for_token: The first 'for' token of the comprehension.
+    has_split_at_for: Whether there is a newline immediately before the
+        for_token.
+    has_interior_split: Whether there is a newline within the comprehension.
+        That is, a split somewhere after expr_token or before closing_bracket.
+  """
+
+  def __init__(self, expr_token):
+    self.expr_token = expr_token
+    self.for_token = None
+    self.has_split_at_for = False
+    self.has_interior_split = False
+
+  def HasTrivialExpr(self):
+    """Returns whether the comp_expr is "trivial" i.e. is a single token."""
+    return self.expr_token.next_token.value == 'for'
+
+  @property
+  def opening_bracket(self):
+    return self.expr_token.previous_token
+
+  @property
+  def closing_bracket(self):
+    return self.opening_bracket.matching_bracket
+
+  def Clone(self):
+    clone = _ComprehensionState(self.expr_token)
+    clone.for_token = self.for_token
+    clone.has_split_at_for = self.has_split_at_for
+    clone.has_interior_split = self.has_interior_split
+    return clone
+
+  def __repr__(self):
+    return ('[opening_bracket::%s, for_token::%s, has_split_at_for::%s,'
+            ' has_interior_split::%s, has_trivial_expr::%s]' %
+            (self.opening_bracket, self.for_token, self.has_split_at_for,
+             self.has_interior_split, self.HasTrivialExpr()))
+
+  def __eq__(self, other):
+    return hash(self) == hash(other)
+
+  def __ne__(self, other):
+    return not self == other
+
+  def __hash__(self, *args, **kwargs):
+    return hash((self.expr_token, self.for_token, self.has_split_at_for,
+                 self.has_interior_split))
diff --git a/yapf/yapflib/format_token.py b/yapf/yapflib/format_token.py
index 8895276..85af621 100644
--- a/yapf/yapflib/format_token.py
+++ b/yapf/yapflib/format_token.py
@@ -49,6 +49,7 @@ class Subtype(object):
   DICTIONARY_KEY_PART = 12
   DICTIONARY_VALUE = 13
   DICT_SET_GENERATOR = 14
+  COMP_EXPR = 21
   COMP_FOR = 15
   COMP_IF = 16
   FUNC_DEF = 17
diff --git a/yapf/yapflib/split_penalty.py b/yapf/yapflib/split_penalty.py
index 55a7134..8ac0780 100644
--- a/yapf/yapflib/split_penalty.py
+++ b/yapf/yapflib/split_penalty.py
@@ -190,8 +190,13 @@ class _SplitPenaltyAssigner(pytree_visitor.PyTreeVisitor):
       if name in {'argument', 'comparison'}:
         # Don't split an argument list with one element if at all possible.
         _SetStronglyConnected(node.children[1])
-        _SetSplitPenalty(
-            _FirstChildNode(node.children[1]), ONE_ELEMENT_ARGUMENT)
+        if (len(node.children[1].children) > 1 and
+            pytree_utils.NodeName(node.children[1].children[1]) == 'comp_for'):
+          # Don't penalize spliting before a comp_for expression.
+          _SetSplitPenalty(_FirstChildNode(node.children[1]), 0)
+        else:
+          _SetSplitPenalty(
+              _FirstChildNode(node.children[1]), ONE_ELEMENT_ARGUMENT)
       elif (pytree_utils.NodeName(node.children[0]) == 'LSQB' and
             len(node.children[1].children) > 2 and
             (name.endswith('_test') or name.endswith('_expr'))):
diff --git a/yapf/yapflib/style.py b/yapf/yapflib/style.py
index 8f30ab0..948c2fd 100644
--- a/yapf/yapflib/style.py
+++ b/yapf/yapflib/style.py
@@ -178,6 +178,22 @@ _STYLE_HELP = dict(
       after."""),
     SPLIT_BEFORE_NAMED_ASSIGNS=textwrap.dedent("""\
       Split named assignments onto individual lines."""),
+    SPLIT_COMPLEX_COMPREHENSION=textwrap.dedent("""\
+      Set to True to split list comprehensions and generators that have
+      non-trivial expressions and multiple clauses before each of these
+      clauses. For example:
+
+        result = [
+            a_long_var + 100 for a_long_var in xrange(1000)
+            if a_long_var % 10]
+
+      would reformat to something like:
+
+        result = [
+            a_long_var + 100
+            for a_long_var in xrange(1000)
+            if a_long_var % 10]
+      """),
     SPLIT_PENALTY_AFTER_OPENING_BRACKET=textwrap.dedent("""\
       The penalty for splitting right after the opening bracket."""),
     SPLIT_PENALTY_AFTER_UNARY_OPERATOR=textwrap.dedent("""\
@@ -187,6 +203,9 @@ _STYLE_HELP = dict(
     SPLIT_PENALTY_BITWISE_OPERATOR=textwrap.dedent("""\
       The penalty of splitting the line around the '&', '|', and '^'
       operators."""),
+    SPLIT_PENALTY_COMPREHENSION=textwrap.dedent("""\
+      The penalty for splitting a list comprehension or generator
+      expression."""),
     SPLIT_PENALTY_EXCESS_CHARACTER=textwrap.dedent("""\
       The penalty for characters over the column limit."""),
     SPLIT_PENALTY_FOR_ADDED_LINE_SPLIT=textwrap.dedent("""\
@@ -243,10 +262,12 @@ def CreatePEP8Style():
       SPLIT_BEFORE_FIRST_ARGUMENT=False,
       SPLIT_BEFORE_LOGICAL_OPERATOR=True,
       SPLIT_BEFORE_NAMED_ASSIGNS=True,
+      SPLIT_COMPLEX_COMPREHENSION=False,
       SPLIT_PENALTY_AFTER_OPENING_BRACKET=30,
       SPLIT_PENALTY_AFTER_UNARY_OPERATOR=10000,
       SPLIT_PENALTY_BEFORE_IF_EXPR=0,
       SPLIT_PENALTY_BITWISE_OPERATOR=300,
+      SPLIT_PENALTY_COMPREHENSION=80,
       SPLIT_PENALTY_EXCESS_CHARACTER=4500,
       SPLIT_PENALTY_FOR_ADDED_LINE_SPLIT=30,
       SPLIT_PENALTY_IMPORT_NAMES=0,
@@ -266,6 +287,8 @@ def CreateGoogleStyle():
   style['SPACE_BETWEEN_ENDING_COMMA_AND_CLOSING_BRACKET'] = False
   style['SPLIT_BEFORE_LOGICAL_OPERATOR'] = False
   style['SPLIT_BEFORE_BITWISE_OPERATOR'] = False
+  style['SPLIT_COMPLEX_COMPREHENSION'] = True
+  style['SPLIT_PENALTY_COMPREHENSION'] = 2100
   return style
 
 
@@ -368,10 +391,12 @@ _STYLE_OPTION_VALUE_CONVERTER = dict(
     SPLIT_BEFORE_FIRST_ARGUMENT=_BoolConverter,
     SPLIT_BEFORE_LOGICAL_OPERATOR=_BoolConverter,
     SPLIT_BEFORE_NAMED_ASSIGNS=_BoolConverter,
+    SPLIT_COMPLEX_COMPREHENSION=_BoolConverter,
     SPLIT_PENALTY_AFTER_OPENING_BRACKET=int,
     SPLIT_PENALTY_AFTER_UNARY_OPERATOR=int,
     SPLIT_PENALTY_BEFORE_IF_EXPR=int,
     SPLIT_PENALTY_BITWISE_OPERATOR=int,
+    SPLIT_PENALTY_COMPREHENSION=int,
     SPLIT_PENALTY_EXCESS_CHARACTER=int,
     SPLIT_PENALTY_FOR_ADDED_LINE_SPLIT=int,
     SPLIT_PENALTY_IMPORT_NAMES=int,
diff --git a/yapf/yapflib/subtype_assigner.py b/yapf/yapflib/subtype_assigner.py
index 213ce31..f24dbbc 100644
--- a/yapf/yapflib/subtype_assigner.py
+++ b/yapf/yapflib/subtype_assigner.py
@@ -304,6 +304,12 @@ class _SubtypeAssigner(pytree_visitor.PyTreeVisitor):
   def Visit_comp_for(self, node):  # pylint: disable=invalid-name
     # comp_for ::= 'for' exprlist 'in' testlist_safe [comp_iter]
     _AppendSubtypeRec(node, format_token.Subtype.COMP_FOR)
+    # Mark the previous node as COMP_EXPR unless this is a nested comprehension
+    # as these will have the outer comprehension as their previous node.
+    attr = pytree_utils.GetNodeAnnotation(node.parent,
+                                          pytree_utils.Annotation.SUBTYPE)
+    if not attr or format_token.Subtype.COMP_FOR not in attr:
+      _AppendSubtypeRec(node.parent.children[0], format_token.Subtype.COMP_EXPR)
     self.DefaultNodeVisit(node)
 
   def Visit_comp_if(self, node):  # pylint: disable=invalid-name


Test patch:
diff --git a/yapftests/pytree_unwrapper_test.py b/yapftests/pytree_unwrapper_test.py
index ab634a5..03a93f0 100644
--- a/yapftests/pytree_unwrapper_test.py
+++ b/yapftests/pytree_unwrapper_test.py
@@ -34,7 +34,8 @@ class PytreeUnwrapperTest(yapf_test_helper.YAPFTest):
     actual = []
     for uwl in uwlines:
       filtered_values = [
-          ft.value for ft in uwl.tokens
+          ft.value
+          for ft in uwl.tokens
           if ft.name not in pytree_utils.NONSEMANTIC_TOKENS
       ]
       actual.append((uwl.depth, filtered_values))
@@ -299,7 +300,8 @@ class MatchBracketsTest(yapf_test_helper.YAPFTest):
     """
     actual = []
     for uwl in uwlines:
-      filtered_values = [(ft, ft.matching_bracket) for ft in uwl.tokens
+      filtered_values = [(ft, ft.matching_bracket)
+                         for ft in uwl.tokens
                          if ft.name not in pytree_utils.NONSEMANTIC_TOKENS]
       if filtered_values:
         actual.append(filtered_values)
diff --git a/yapftests/reformatter_basic_test.py b/yapftests/reformatter_basic_test.py
index cef399a..aa237e1 100644
--- a/yapftests/reformatter_basic_test.py
+++ b/yapftests/reformatter_basic_test.py
@@ -457,6 +457,74 @@ class BasicReformatterTest(yapf_test_helper.YAPFTest):
     uwlines = yapf_test_helper.ParseAndUnwrap(unformatted_code)
     self.assertCodeEqual(expected_formatted_code, reformatter.Reformat(uwlines))
 
+  def testListComprehensionPreferOneLine(self):
+    unformatted_code = textwrap.dedent("""\
+        def given(y):
+            long_variable_name = [
+                long_var_name + 1
+                for long_var_name in ()
+                if long_var_name == 2]
+        """)
+    expected_formatted_code = textwrap.dedent("""\
+        def given(y):
+          long_variable_name = [
+              long_var_name + 1 for long_var_name in () if long_var_name == 2
+          ]
+        """)
+    uwlines = yapf_test_helper.ParseAndUnwrap(unformatted_code)
+    self.assertCodeEqual(expected_formatted_code, reformatter.Reformat(uwlines))
+
+  def testListComprehensionPreferOneLineOverArithmeticSplit(self):
+    unformatted_code = textwrap.dedent("""\
+        def given(used_identifiers):
+          return (sum(len(identifier)
+                      for identifier in used_identifiers) / len(used_identifiers))
+        """)
+    expected_formatted_code = textwrap.dedent("""\
+        def given(used_identifiers):
+          return (sum(len(identifier) for identifier in used_identifiers) /
+                  len(used_identifiers))
+        """)
+    uwlines = yapf_test_helper.ParseAndUnwrap(unformatted_code)
+    self.assertCodeEqual(expected_formatted_code, reformatter.Reformat(uwlines))
+
+  def testListComprehensionPreferThreeLinesForLineWrap(self):
+    unformatted_code = textwrap.dedent("""\
+        def given(y):
+            long_variable_name = [
+                long_var_name + 1
+                for long_var_name, number_two in ()
+                if long_var_name == 2 and number_two == 3]
+        """)
+    expected_formatted_code = textwrap.dedent("""\
+        def given(y):
+          long_variable_name = [
+              long_var_name + 1
+              for long_var_name, number_two in ()
+              if long_var_name == 2 and number_two == 3
+          ]
+        """)
+    uwlines = yapf_test_helper.ParseAndUnwrap(unformatted_code)
+    self.assertCodeEqual(expected_formatted_code, reformatter.Reformat(uwlines))
+
+  def testListComprehensionPreferNoBreakForTrivialExpression(self):
+    unformatted_code = textwrap.dedent("""\
+        def given(y):
+            long_variable_name = [
+                long_var_name
+                for long_var_name, number_two in ()
+                if long_var_name == 2 and number_two == 3]
+        """)
+    expected_formatted_code = textwrap.dedent("""\
+        def given(y):
+          long_variable_name = [
+              long_var_name for long_var_name, number_two in ()
+              if long_var_name == 2 and number_two == 3
+          ]
+        """)
+    uwlines = yapf_test_helper.ParseAndUnwrap(unformatted_code)
+    self.assertCodeEqual(expected_formatted_code, reformatter.Reformat(uwlines))
+
   def testOpeningAndClosingBrackets(self):
     unformatted_code = """\
 foo( (1, ) )
@@ -494,8 +562,8 @@ foo((
 find_symbol(node.type) + "< " + " ".join(find_pattern(n) for n in node.child) + " >"
 """
     expected_formatted_code = """\
-find_symbol(node.type) + "< " + " ".join(find_pattern(n)
-                                         for n in node.child) + " >"
+find_symbol(node.type) + "< " + " ".join(
+    find_pattern(n) for n in node.child) + " >"
 """
     uwlines = yapf_test_helper.ParseAndUnwrap(unformatted_code)
     self.assertCodeEqual(expected_formatted_code, reformatter.Reformat(uwlines))
@@ -887,7 +955,8 @@ class foo:
           if True:
             if True:
               python_files.extend(
-                  os.path.join(filename, f) for f in os.listdir(filename)
+                  os.path.join(filename, f)
+                  for f in os.listdir(filename)
                   if IsPythonFile(os.path.join(filename, f)))
         """)
     uwlines = yapf_test_helper.ParseAndUnwrap(code)
@@ -1221,7 +1290,8 @@ s = 'foo \\
     code = textwrap.dedent("""\
         foo = {
             variable: 'hello world. How are you today?'
-            for variable in fnord if variable != 37
+            for variable in fnord
+            if variable != 37
         }
         """)
     uwlines = yapf_test_helper.ParseAndUnwrap(code)
diff --git a/yapftests/reformatter_buganizer_test.py b/yapftests/reformatter_buganizer_test.py
index db378b6..ae9be47 100644
--- a/yapftests/reformatter_buganizer_test.py
+++ b/yapftests/reformatter_buganizer_test.py
@@ -279,8 +279,8 @@ checkpoint_files = gfile.Glob(
 
   def testB37460004(self):
     code = textwrap.dedent("""\
-        assert all(s not in (_SENTINEL, None) for s in
-                   nested_schemas), 'Nested schemas should never contain None/_SENTINEL'
+        assert all(s not in (_SENTINEL, None) for s in nested_schemas
+                  ), 'Nested schemas should never contain None/_SENTINEL'
         """)
     uwlines = yapf_test_helper.ParseAndUnwrap(code)
     self.assertCodeEqual(code, reformatter.Reformat(uwlines))
diff --git a/yapftests/split_penalty_test.py b/yapftests/split_penalty_test.py
index deaf444..1c8bf4a 100644
--- a/yapftests/split_penalty_test.py
+++ b/yapftests/split_penalty_test.py
@@ -237,6 +237,23 @@ class SplitPenaltyTest(unittest.TestCase):
         (')', VERY_STRONGLY_CONNECTED),
     ])
 
+    # Test single generator argument.
+    code = 'max(i for i in xrange(10))\n'
+    tree = self._ParseAndComputePenalties(code)
+    self._CheckPenalties(tree, [
+        ('max', None),
+        ('(', UNBREAKABLE),
+        ('i', 0),
+        ('for', 0),
+        ('i', STRONGLY_CONNECTED),
+        ('in', STRONGLY_CONNECTED),
+        ('xrange', STRONGLY_CONNECTED),
+        ('(', UNBREAKABLE),
+        ('10', STRONGLY_CONNECTED),
+        (')', VERY_STRONGLY_CONNECTED),
+        (')', VERY_STRONGLY_CONNECTED),
+    ])
+
 
 if __name__ == '__main__':
   unittest.main()
diff --git a/yapftests/subtype_assigner_test.py b/yapftests/subtype_assigner_test.py
index 766e8de..9dd1050 100644
--- a/yapftests/subtype_assigner_test.py
+++ b/yapftests/subtype_assigner_test.py
@@ -34,7 +34,8 @@ class SubtypeAssignerTest(yapf_test_helper.YAPFTest):
     """
     actual = []
     for uwl in uwlines:
-      filtered_values = [(ft.value, ft.subtypes) for ft in uwl.tokens
+      filtered_values = [(ft.value, ft.subtypes)
+                         for ft in uwl.tokens
                          if ft.name not in pytree_utils.NONSEMANTIC_TOKENS]
       if filtered_values:
         actual.append(filtered_values)
@@ -111,11 +112,11 @@ class SubtypeAssignerTest(yapf_test_helper.YAPFTest):
          (':', [format_token.Subtype.NONE])],
         [('return', [format_token.Subtype.NONE]),
          ('{', [format_token.Subtype.NONE]),
-         ('s', [format_token.Subtype.NONE]),
-         ('.', [format_token.Subtype.NONE]),
-         ('lower', [format_token.Subtype.NONE]),
-         ('(', [format_token.Subtype.NONE]),
-         (')', [format_token.Subtype.NONE]),
+         ('s', {format_token.Subtype.COMP_EXPR}),
+         ('.', {format_token.Subtype.COMP_EXPR}),
+         ('lower', {format_token.Subtype.COMP_EXPR}),
+         ('(', {format_token.Subtype.COMP_EXPR}),
+         (')', {format_token.Subtype.COMP_EXPR}),
          ('for', {format_token.Subtype.DICT_SET_GENERATOR,
                   format_token.Subtype.COMP_FOR}),
          ('s', {format_token.Subtype.COMP_FOR}),