Repository: google/yapf
Pull request #1044 (2022-12-26 16:15:07)

Generated dicts also have keys and values for rules to apply
The general idea:

If
```python
responseDict = {'childDict': {'spam': 'eggs'}}
```
is formatted as 
```python
responseDict = {
    'childDict': {
        'spam': 'eggs'
    }
}
```
so should
```python
generatedDict = {x: x for x in 'value'}
```
be formatted as 
```python
generatedDict = {
    x: x for x in 'value'
}
```
assuming `force_multiline_dict` is enabled. Also allows to achieve the following format if `split_before_dict_set_generator` is enabled:
```python
generatedDict = {
    x: x
    for x in 'value'
}
```

bwendling: Thanks! Please add a comment to CHANGELOG.
alexey-pelykh: @bwendling not sure that done correctly, yet done nonetheless

Patch:
diff --git a/CHANGELOG b/CHANGELOG
index 8b29032..f3ba6e9 100644
--- a/CHANGELOG
+++ b/CHANGELOG
@@ -9,6 +9,7 @@
 ### Changes
 - Moved 'pytree' parsing tools into its own subdirectory.
 - Add support for Python 3.10.
+- Format generated dicts with respect to same rules as regular dicts
 ### Fixed
 - Split line before all comparison operators.
 
diff --git a/yapf/pytree/subtype_assigner.py b/yapf/pytree/subtype_assigner.py
index dd3ea3d..c93d69d 100644
--- a/yapf/pytree/subtype_assigner.py
+++ b/yapf/pytree/subtype_assigner.py
@@ -66,20 +66,26 @@ class _SubtypeAssigner(pytree_visitor.PyTreeVisitor):
     for child in node.children:
       self.Visit(child)
 
-    comp_for = False
     dict_maker = False
 
+    def markAsDictSetGenerator(node):
+      _AppendFirstLeafTokenSubtype(node, subtypes.DICT_SET_GENERATOR)
+      for child in node.children:
+        if pytree_utils.NodeName(child) == 'comp_for':
+          markAsDictSetGenerator(child)
+
     for child in node.children:
       if pytree_utils.NodeName(child) == 'comp_for':
-        comp_for = True
-        _AppendFirstLeafTokenSubtype(child, subtypes.DICT_SET_GENERATOR)
+        markAsDictSetGenerator(child)
       elif child.type in (grammar_token.COLON, grammar_token.DOUBLESTAR):
         dict_maker = True
 
-    if not comp_for and dict_maker:
+    if dict_maker:
       last_was_colon = False
       unpacking = False
       for child in node.children:
+        if pytree_utils.NodeName(child) == 'comp_for':
+          break
         if child.type == grammar_token.DOUBLESTAR:
           _AppendFirstLeafTokenSubtype(child, subtypes.KWARGS_STAR_STAR)
         if last_was_colon:
@@ -335,7 +341,10 @@ class _SubtypeAssigner(pytree_visitor.PyTreeVisitor):
     attr = pytree_utils.GetNodeAnnotation(node.parent,
                                           pytree_utils.Annotation.SUBTYPE)
     if not attr or subtypes.COMP_FOR not in attr:
-      _AppendSubtypeRec(node.parent.children[0], subtypes.COMP_EXPR)
+      sibling = node.prev_sibling
+      while sibling:
+        _AppendSubtypeRec(sibling, subtypes.COMP_EXPR)
+        sibling = sibling.prev_sibling
     self.DefaultNodeVisit(node)
 
   def Visit_old_comp_for(self, node):  # pylint: disable=invalid-name
diff --git a/yapf/yapflib/format_decision_state.py b/yapf/yapflib/format_decision_state.py
index c299d1c..607f8c5 100644
--- a/yapf/yapflib/format_decision_state.py
+++ b/yapf/yapflib/format_decision_state.py
@@ -1031,6 +1031,8 @@ class FormatDecisionState(object):
     current = opening.next_token.next_token
 
     while current and current != closing:
+      if subtypes.DICT_SET_GENERATOR in current.subtypes:
+        break
       if subtypes.DICTIONARY_KEY in current.subtypes:
         prev = PreviousNonCommentToken(current)
         if prev.value == ',':


Test patch:
diff --git a/yapftests/reformatter_basic_test.py b/yapftests/reformatter_basic_test.py
index 657d1e2..935c7c3 100644
--- a/yapftests/reformatter_basic_test.py
+++ b/yapftests/reformatter_basic_test.py
@@ -1624,6 +1624,18 @@ s = 'foo \\
     llines = yapf_test_helper.ParseAndUnwrap(code)
     self.assertCodeEqual(code, reformatter.Reformat(llines))
 
+    unformatted_code = textwrap.dedent("""\
+        foo = {
+            x: x
+            for x in fnord
+        }
+        """)  # noqa
+    expected_code = textwrap.dedent("""\
+        foo = {x: x for x in fnord}
+        """)
+    llines = yapf_test_helper.ParseAndUnwrap(unformatted_code)
+    self.assertCodeEqual(expected_code, reformatter.Reformat(llines))
+
   def testUnaryOpInDictionaryValue(self):
     code = textwrap.dedent("""\
         beta = "123"
@@ -3123,8 +3135,10 @@ my_dict = {
     try:
       style.SetGlobalStyle(
           style.CreateStyleFromConfig('{force_multiline_dict: true}'))
-      unformatted_code = textwrap.dedent(
-          "responseDict = {'childDict': {'spam': 'eggs'}}\n")
+      unformatted_code = textwrap.dedent("""\
+        responseDict = {'childDict': {'spam': 'eggs'}}
+        generatedDict = {x: x for x in 'value'}
+      """)
       llines = yapf_test_helper.ParseAndUnwrap(unformatted_code)
       actual = reformatter.Reformat(llines)
       expected = textwrap.dedent("""\
@@ -3133,6 +3147,9 @@ my_dict = {
                 'spam': 'eggs'
             }
         }
+        generatedDict = {
+            x: x for x in 'value'
+        }
       """)
       self.assertCodeEqual(expected, actual)
     finally:
@@ -3144,6 +3161,7 @@ my_dict = {
           style.CreateStyleFromConfig('{force_multiline_dict: false}'))
       unformatted_code = textwrap.dedent("""\
         responseDict = {'childDict': {'spam': 'eggs'}}
+        generatedDict = {x: x for x in 'value'}
       """)
       expected_formatted_code = unformatted_code
       llines = yapf_test_helper.ParseAndUnwrap(unformatted_code)
diff --git a/yapftests/subtype_assigner_test.py b/yapftests/subtype_assigner_test.py
index 8616169..c705077 100644
--- a/yapftests/subtype_assigner_test.py
+++ b/yapftests/subtype_assigner_test.py
@@ -130,6 +130,36 @@ class SubtypeAssignerTest(yapf_test_helper.YAPFTest):
     ])
 
   def testSetComprehension(self):
+    code = textwrap.dedent("""\
+        def foo(value):
+          return {value.lower()}
+        """)
+    llines = yapf_test_helper.ParseAndUnwrap(code)
+    self._CheckFormatTokenSubtypes(llines, [
+        [
+            ('def', {subtypes.NONE}),
+            ('foo', {subtypes.FUNC_DEF}),
+            ('(', {subtypes.NONE}),
+            ('value', {
+                subtypes.NONE,
+                subtypes.PARAMETER_START,
+                subtypes.PARAMETER_STOP,
+            }),
+            (')', {subtypes.NONE}),
+            (':', {subtypes.NONE}),
+        ],
+        [
+            ('return', {subtypes.NONE}),
+            ('{', {subtypes.NONE}),
+            ('value', {subtypes.NONE}),
+            ('.', {subtypes.NONE}),
+            ('lower', {subtypes.NONE}),
+            ('(', {subtypes.NONE}),
+            (')', {subtypes.NONE}),
+            ('}', {subtypes.NONE}),
+        ],
+    ])
+
     code = textwrap.dedent("""\
         def foo(strs):
           return {s.lower() for s in strs}
@@ -167,6 +197,209 @@ class SubtypeAssignerTest(yapf_test_helper.YAPFTest):
         ],
     ])
 
+    code = textwrap.dedent("""\
+        def foo(strs):
+          return {s + s.lower() for s in strs}
+        """)
+    llines = yapf_test_helper.ParseAndUnwrap(code)
+    self._CheckFormatTokenSubtypes(llines, [
+        [
+            ('def', {subtypes.NONE}),
+            ('foo', {subtypes.FUNC_DEF}),
+            ('(', {subtypes.NONE}),
+            ('strs', {
+                subtypes.NONE,
+                subtypes.PARAMETER_START,
+                subtypes.PARAMETER_STOP,
+            }),
+            (')', {subtypes.NONE}),
+            (':', {subtypes.NONE}),
+        ],
+        [
+            ('return', {subtypes.NONE}),
+            ('{', {subtypes.NONE}),
+            ('s', {subtypes.COMP_EXPR}),
+            ('+', {subtypes.BINARY_OPERATOR, subtypes.COMP_EXPR}),
+            ('s', {subtypes.COMP_EXPR}),
+            ('.', {subtypes.COMP_EXPR}),
+            ('lower', {subtypes.COMP_EXPR}),
+            ('(', {subtypes.COMP_EXPR}),
+            (')', {subtypes.COMP_EXPR}),
+            ('for', {
+                subtypes.DICT_SET_GENERATOR,
+                subtypes.COMP_FOR,
+            }),
+            ('s', {subtypes.COMP_FOR}),
+            ('in', {subtypes.COMP_FOR}),
+            ('strs', {subtypes.COMP_FOR}),
+            ('}', {subtypes.NONE}),
+        ],
+    ])
+
+    code = textwrap.dedent("""\
+        def foo(strs):
+          return {c.lower() for s in strs for c in s}
+        """)
+    llines = yapf_test_helper.ParseAndUnwrap(code)
+    self._CheckFormatTokenSubtypes(llines, [
+        [
+            ('def', {subtypes.NONE}),
+            ('foo', {subtypes.FUNC_DEF}),
+            ('(', {subtypes.NONE}),
+            ('strs', {
+                subtypes.NONE,
+                subtypes.PARAMETER_START,
+                subtypes.PARAMETER_STOP,
+            }),
+            (')', {subtypes.NONE}),
+            (':', {subtypes.NONE}),
+        ],
+        [
+            ('return', {subtypes.NONE}),
+            ('{', {subtypes.NONE}),
+            ('c', {subtypes.COMP_EXPR}),
+            ('.', {subtypes.COMP_EXPR}),
+            ('lower', {subtypes.COMP_EXPR}),
+            ('(', {subtypes.COMP_EXPR}),
+            (')', {subtypes.COMP_EXPR}),
+            ('for', {
+                subtypes.DICT_SET_GENERATOR,
+                subtypes.COMP_FOR,
+                subtypes.COMP_EXPR,
+            }),
+            ('s', {subtypes.COMP_FOR, subtypes.COMP_EXPR}),
+            ('in', {subtypes.COMP_FOR, subtypes.COMP_EXPR}),
+            ('strs', {subtypes.COMP_FOR, subtypes.COMP_EXPR}),
+            ('for', {
+                subtypes.DICT_SET_GENERATOR,
+                subtypes.COMP_FOR,
+            }),
+            ('c', {subtypes.COMP_FOR}),
+            ('in', {subtypes.COMP_FOR}),
+            ('s', {subtypes.COMP_FOR}),
+            ('}', {subtypes.NONE}),
+        ],
+    ])
+
+  def testDictComprehension(self):
+    code = textwrap.dedent("""\
+        def foo(value):
+          return {value: value.lower()}
+        """)
+    llines = yapf_test_helper.ParseAndUnwrap(code)
+    self._CheckFormatTokenSubtypes(llines, [
+        [
+            ('def', {subtypes.NONE}),
+            ('foo', {subtypes.FUNC_DEF}),
+            ('(', {subtypes.NONE}),
+            ('value', {
+                subtypes.NONE,
+                subtypes.PARAMETER_START,
+                subtypes.PARAMETER_STOP,
+            }),
+            (')', {subtypes.NONE}),
+            (':', {subtypes.NONE}),
+        ],
+        [
+            ('return', {subtypes.NONE}),
+            ('{', {subtypes.NONE}),
+            ('value', {subtypes.DICTIONARY_KEY, subtypes.DICTIONARY_KEY_PART}),
+            (':', {subtypes.NONE}),
+            ('value', {subtypes.DICTIONARY_VALUE}),
+            ('.', {subtypes.NONE}),
+            ('lower', {subtypes.NONE}),
+            ('(', {subtypes.NONE}),
+            (')', {subtypes.NONE}),
+            ('}', {subtypes.NONE}),
+        ],
+    ])
+
+    code = textwrap.dedent("""\
+        def foo(strs):
+          return {s: s.lower() for s in strs}
+        """)
+    llines = yapf_test_helper.ParseAndUnwrap(code)
+    self._CheckFormatTokenSubtypes(llines, [
+        [
+            ('def', {subtypes.NONE}),
+            ('foo', {subtypes.FUNC_DEF}),
+            ('(', {subtypes.NONE}),
+            ('strs', {
+                subtypes.NONE,
+                subtypes.PARAMETER_START,
+                subtypes.PARAMETER_STOP,
+            }),
+            (')', {subtypes.NONE}),
+            (':', {subtypes.NONE}),
+        ],
+        [
+            ('return', {subtypes.NONE}),
+            ('{', {subtypes.NONE}),
+            ('s', {subtypes.DICTIONARY_KEY, subtypes.DICTIONARY_KEY_PART, subtypes.COMP_EXPR}),
+            (':', {subtypes.COMP_EXPR}),
+            ('s', {subtypes.DICTIONARY_VALUE, subtypes.COMP_EXPR}),
+            ('.', {subtypes.COMP_EXPR}),
+            ('lower', {subtypes.COMP_EXPR}),
+            ('(', {subtypes.COMP_EXPR}),
+            (')', {subtypes.COMP_EXPR}),
+            ('for', {
+                subtypes.DICT_SET_GENERATOR,
+                subtypes.COMP_FOR,
+            }),
+            ('s', {subtypes.COMP_FOR}),
+            ('in', {subtypes.COMP_FOR}),
+            ('strs', {subtypes.COMP_FOR}),
+            ('}', {subtypes.NONE}),
+        ],
+    ])
+
+    code = textwrap.dedent("""\
+        def foo(strs):
+          return {c: c.lower() for s in strs for c in s}
+        """)
+    llines = yapf_test_helper.ParseAndUnwrap(code)
+    self._CheckFormatTokenSubtypes(llines, [
+        [
+            ('def', {subtypes.NONE}),
+            ('foo', {subtypes.FUNC_DEF}),
+            ('(', {subtypes.NONE}),
+            ('strs', {
+                subtypes.NONE,
+                subtypes.PARAMETER_START,
+                subtypes.PARAMETER_STOP,
+            }),
+            (')', {subtypes.NONE}),
+            (':', {subtypes.NONE}),
+        ],
+        [
+            ('return', {subtypes.NONE}),
+            ('{', {subtypes.NONE}),
+            ('c', {subtypes.DICTIONARY_KEY, subtypes.DICTIONARY_KEY_PART, subtypes.COMP_EXPR}),
+            (':', {subtypes.COMP_EXPR}),
+            ('c', {subtypes.DICTIONARY_VALUE, subtypes.COMP_EXPR}),
+            ('.', {subtypes.COMP_EXPR}),
+            ('lower', {subtypes.COMP_EXPR}),
+            ('(', {subtypes.COMP_EXPR}),
+            (')', {subtypes.COMP_EXPR}),
+            ('for', {
+                subtypes.DICT_SET_GENERATOR,
+                subtypes.COMP_FOR,
+                subtypes.COMP_EXPR,
+            }),
+            ('s', {subtypes.COMP_FOR, subtypes.COMP_EXPR}),
+            ('in', {subtypes.COMP_FOR, subtypes.COMP_EXPR}),
+            ('strs', {subtypes.COMP_FOR, subtypes.COMP_EXPR}),
+            ('for', {
+                subtypes.DICT_SET_GENERATOR,
+                subtypes.COMP_FOR,
+            }),
+            ('c', {subtypes.COMP_FOR}),
+            ('in', {subtypes.COMP_FOR}),
+            ('s', {subtypes.COMP_FOR}),
+            ('}', {subtypes.NONE}),
+        ],
+    ])
+
   def testUnaryNotOperator(self):
     code = textwrap.dedent("""\
         not a