Wrap the in clause of comprehensions across lines if necessary (#4699)

* CI: Remove now-uneeded workarounds The issues mentioned in the comments have all been fixed Signed-off-by: cobalt <61329810+cobaltt7@users.noreply.github.com> * Fix Click typing issue in run self * oops, run self * i swear this was failing earlier ig it's not that related to this pr anyway * has click updated itself yet on pre-commit.ci * Bump click Signed-off-by: cobalt <61329810+cobaltt7@users.noreply.github.com> * Split the `in` clause of comprehensions onto its own line if necessary Fixes #3498 Signed-off-by: cobalt <61329810+cobaltt7@users.noreply.github.com> * Lint issues Signed-off-by: cobalt <61329810+cobaltt7@users.noreply.github.com> * Move to preview style Signed-off-by: cobalt <61329810+cobaltt7@users.noreply.github.com> * run self Signed-off-by: cobalt <61329810+cobaltt7@users.noreply.github.com> * Alternative approach Signed-off-by: cobalt <61329810+cobaltt7@users.noreply.github.com> * oops! Signed-off-by: cobalt <61329810+cobaltt7@users.noreply.github.com> * run test on preview, not unstable Signed-off-by: cobalt <61329810+cobaltt7@users.noreply.github.com> * Don't remove parens around ternaries Signed-off-by: cobalt <61329810+cobaltt7@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update src/black/linegen.py * Update docs/the_black_code_style/future_style.md --------- Signed-off-by: cobalt <61329810+cobaltt7@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: cobaltt7

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
+}