Wrap the in clause of comprehensions across lines if necessary

CHANGES.md

@@ -34,6 +34,7 @@
 - Improve `multiline_string_handling` with ternaries and dictionaries (#4657)
 - Fix a bug where `string_processing` would not split f-strings directly after
   expressions (#4680)
+- Wrap the `in` clause of comprehensions across lines if necessary (#4699)
 - Remove parentheses around multiple exception types in `except` and `except*` without
   `as`. (#4720)
 

docs/the_black_code_style/future_style.md

@@ -27,8 +27,10 @@ Currently, the following features are included in the preview style:
 - `wrap_long_dict_values_in_parens`: Add parentheses around long values in dictionaries
   ([see below](labels/wrap-long-dict-values))
 - `fix_fmt_skip_in_one_liners`: Fix `# fmt: skip` behaviour on one-liner declarations,
-  such as `def foo(): return "mock"  # fmt: skip`, where previously the declaration
-  would have been incorrectly collapsed.
+  such as `def foo(): return "mock" # fmt: skip`, where previously the declaration would
+  have been incorrectly collapsed.
+- `wrap_comprehension_in`: Wrap the `in` clause of list and dictionary comprehensions
+  across lines if it would otherwise exceed the maximum line length.
 - `remove_parens_around_except_types`: Remove parentheses around multiple exception
   types in `except` and `except*` without `as`. See PEP 758 for details.
 

src/black/linegen.py

@@ -579,6 +579,16 @@ def visit_fstring(self, node: Node) -> Iterator[Line]:
 
         # yield from self.visit_default(node)
 
+    def visit_comp_for(self, node: Node) -> Iterator[Line]:
+        if Preview.wrap_comprehension_in in self.mode:
+            normalize_invisible_parens(
+                node, parens_after={"in"}, mode=self.mode, features=self.features
+            )
+        yield from self.visit_default(node)
+
+    def visit_old_comp_for(self, node: Node) -> Iterator[Line]:
+        yield from self.visit_comp_for(node)
+
     def __post_init__(self) -> None:
         """You are in a twisty little maze of passages."""
         self.current_line = Line(mode=self.mode)
@@ -1466,7 +1476,13 @@ def normalize_invisible_parens(  # noqa: C901
                     wrap_in_parentheses(node, child, visible=False)
             elif isinstance(child, Node) and node.type == syms.with_stmt:
                 remove_with_parens(child, node, mode=mode, features=features)
-            elif child.type == syms.atom:
+            elif child.type == syms.atom and not (
+                "in" in parens_after
+                and len(child.children) == 3
+                and is_lpar_token(child.children[0])
+                and is_rpar_token(child.children[-1])
+                and child.children[1].type == syms.test
+            ):
                 if maybe_make_parens_invisible_in_atom(
                     child, parent=node, mode=mode, features=features
                 ):

src/black/mode.py

@@ -231,6 +231,7 @@ class Preview(Enum):
     multiline_string_handling = auto()
     always_one_newline_after_import = auto()
     fix_fmt_skip_in_one_liners = auto()
+    wrap_comprehension_in = auto()
     # Remove parentheses around multiple exception types in except and
     # except* without as. See PEP 758 for details.
     remove_parens_around_except_types = auto()

src/black/resources/black.schema.json

@@ -86,6 +86,7 @@
           "multiline_string_handling",
           "always_one_newline_after_import",
           "fix_fmt_skip_in_one_liners",
+          "wrap_comprehension_in",
           "remove_parens_around_except_types"
         ]
       },

tests/data/cases/preview_wrap_comprehension_in.py

@@ -0,0 +1,161 @@
+# flags: --preview --line-length=79
+
+[a for graph_path_expression in refined_constraint.condition_as_predicate.variables]
+[
+    a
+    for graph_path_expression in refined_constraint.condition_as_predicate.variables
+]
+[
+    a
+    for graph_path_expression
+    in refined_constraint.condition_as_predicate.variables
+]
+[
+    a
+    for graph_path_expression in (
+        refined_constraint.condition_as_predicate.variables
+    )
+]
+
+[
+    (foobar_very_long_key, foobar_very_long_value)
+    for foobar_very_long_key, foobar_very_long_value in foobar_very_long_dictionary.items()
+]
+
+# Don't split the `in` if it's not too long
+lcomp3 = [
+    element.split("\n", 1)[0]
+    for element in collection.select_elements()
+    # right
+    if element is not None
+]
+
+# Don't remove parens around ternaries
+expected = [i for i in (a if b else c)]
+
+# Nested arrays
+# First in will not be split because it would still be too long
+[[
+    x
+    for x in bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb
+    for y in xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
+]]
+
+# Multiple comprehensions, only split the second `in`
+graph_path_expressions_in_local_constraint_refinements = [
+    graph_path_expression
+    for refined_constraint in self._local_constraint_refinements.values()
+    if refined_constraint is not None
+    for graph_path_expression in refined_constraint.condition_as_predicate.variables
+]
+
+# Dictionary comprehensions
+dict_with_really_long_names = {
+    really_really_long_key_name: an_even_longer_really_really_long_key_value
+    for really_really_long_key_name, an_even_longer_really_really_long_key_value in really_really_really_long_dict_name.items()
+}
+{
+    key_with_super_really_long_name: key_with_super_really_long_name
+    for key_with_super_really_long_name in dictionary_with_super_really_long_name
+}
+{
+    key_with_super_really_long_name: key_with_super_really_long_name
+    for key_with_super_really_long_name
+    in dictionary_with_super_really_long_name
+}
+{
+    key_with_super_really_long_name: key_with_super_really_long_name
+    for key in (
+        dictionary
+    )
+}
+
+# output
+[
+    a
+    for graph_path_expression in (
+        refined_constraint.condition_as_predicate.variables
+    )
+]
+[
+    a
+    for graph_path_expression in (
+        refined_constraint.condition_as_predicate.variables
+    )
+]
+[
+    a
+    for graph_path_expression in (
+        refined_constraint.condition_as_predicate.variables
+    )
+]
+[
+    a
+    for graph_path_expression in (
+        refined_constraint.condition_as_predicate.variables
+    )
+]
+
+[
+    (foobar_very_long_key, foobar_very_long_value)
+    for foobar_very_long_key, foobar_very_long_value in (
+        foobar_very_long_dictionary.items()
+    )
+]
+
+# Don't split the `in` if it's not too long
+lcomp3 = [
+    element.split("\n", 1)[0]
+    for element in collection.select_elements()
+    # right
+    if element is not None
+]
+
+# Don't remove parens around ternaries
+expected = [i for i in (a if b else c)]
+
+# Nested arrays
+# First in will not be split because it would still be too long
+[
+    [
+        x
+        for x in bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb
+        for y in (
+            xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
+        )
+    ]
+]
+
+# Multiple comprehensions, only split the second `in`
+graph_path_expressions_in_local_constraint_refinements = [
+    graph_path_expression
+    for refined_constraint in self._local_constraint_refinements.values()
+    if refined_constraint is not None
+    for graph_path_expression in (
+        refined_constraint.condition_as_predicate.variables
+    )
+]
+
+# Dictionary comprehensions
+dict_with_really_long_names = {
+    really_really_long_key_name: an_even_longer_really_really_long_key_value
+    for really_really_long_key_name, an_even_longer_really_really_long_key_value in (
+        really_really_really_long_dict_name.items()
+    )
+}
+{
+    key_with_super_really_long_name: key_with_super_really_long_name
+    for key_with_super_really_long_name in (
+        dictionary_with_super_really_long_name
+    )
+}
+{
+    key_with_super_really_long_name: key_with_super_really_long_name
+    for key_with_super_really_long_name in (
+        dictionary_with_super_really_long_name
+    )
+}
+{
+    key_with_super_really_long_name: key_with_super_really_long_name
+    for key in dictionary
+}