diff --git a/CHANGES.md b/CHANGES.md index cc23ec4d9ad..7953d0dc626 100644 --- a/CHANGES.md +++ b/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) diff --git a/docs/the_black_code_style/future_style.md b/docs/the_black_code_style/future_style.md index 45ec1e67ed8..13bcaa94e5d 100644 --- a/docs/the_black_code_style/future_style.md +++ b/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. diff --git a/src/black/linegen.py b/src/black/linegen.py index 49fab818a6d..27c2c92d9b2 100644 --- a/src/black/linegen.py +++ b/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 ): diff --git a/src/black/mode.py b/src/black/mode.py index 86e0bfcb1c2..4d85358d5c5 100644 --- a/src/black/mode.py +++ b/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() diff --git a/src/black/resources/black.schema.json b/src/black/resources/black.schema.json index b342f2d2204..c3d7d03d4cc 100644 --- a/src/black/resources/black.schema.json +++ b/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" ] }, diff --git a/tests/data/cases/preview_wrap_comprehension_in.py b/tests/data/cases/preview_wrap_comprehension_in.py new file mode 100644 index 00000000000..e457f0e772f --- /dev/null +++ b/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 +}