diff --git a/CHANGES.md b/CHANGES.md index b67056766de..c07f1f44559 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -31,6 +31,8 @@ - Fix `fix_fmt_skip_in_one_liners` crashing on `with` statements (#4853) - Fix `fix_fmt_skip_in_one_liners` crashing on annotated parameters (#4854) - Fix `# fmt: skip` behavior for deeply nested expressions (#4883) +- Keep concatenated list comprehensions on their own lines when long additions need to + wrap (#4887) ### Configuration diff --git a/docs/the_black_code_style/future_style.md b/docs/the_black_code_style/future_style.md index 07bc5258d92..7a1ac08f024 100644 --- a/docs/the_black_code_style/future_style.md +++ b/docs/the_black_code_style/future_style.md @@ -24,6 +24,8 @@ Currently, the following features are included in the preview style: between `#` and `type:` or between `type:` and value to `# type: (value)` - `wrap_comprehension_in`: Wrap the `in` clause of list and dictionary comprehensions across lines if it would otherwise exceed the maximum line length. +- `concatenated_list_comprehensions`: Prefer splitting concatenated list comprehensions + between operands, keeping the comprehensions themselves on a single line. - `remove_parens_around_except_types`: Remove parentheses around multiple exception types in `except` and `except*` without `as`. See PEP 758 for details. - `normalize_cr_newlines`: Add `\r` style newlines to the potential newlines to diff --git a/src/black/linegen.py b/src/black/linegen.py index 8d6fa30c49d..af7958f8063 100644 --- a/src/black/linegen.py +++ b/src/black/linegen.py @@ -50,6 +50,7 @@ is_arith_like, is_async_stmt_or_funcdef, is_atom_with_invisible_parens, + is_concatenated_list_comprehension, is_docstring, is_empty_tuple, is_generator, @@ -1041,6 +1042,20 @@ def _first_right_hand_split( return RHSResult(head, body, tail, opening_bracket, closing_bracket) +def _optional_parens_wrap_concatenated_list_comp(opening_bracket: Leaf) -> bool: + parent = opening_bracket.parent + if ( + isinstance(parent, Node) + and parent.type == syms.atom + and len(parent.children) >= 3 + and parent.children[0] is opening_bracket + ): + middle = parent.children[1] + return isinstance(middle, Node) and is_concatenated_list_comprehension(middle) + + return False + + def _maybe_split_omitting_optional_parens( rhs: RHSResult, line: Line, @@ -1048,6 +1063,15 @@ def _maybe_split_omitting_optional_parens( features: Collection[Feature] = (), omit: Collection[LeafID] = (), ) -> Iterator[Line]: + force_optional_parens = ( + Preview.concatenated_list_comprehensions in mode + and rhs.opening_bracket.type == token.LPAR + and not rhs.opening_bracket.value + and rhs.closing_bracket.type == token.RPAR + and not rhs.closing_bracket.value + and _optional_parens_wrap_concatenated_list_comp(rhs.opening_bracket) + ) + if ( Feature.FORCE_OPTIONAL_PARENTHESES not in features # the opening bracket is an optional paren @@ -1061,6 +1085,7 @@ def _maybe_split_omitting_optional_parens( and not line.is_import # and we can actually remove the parens and can_omit_invisible_parens(rhs, mode.line_length) + and not force_optional_parens ): omit = {id(rhs.closing_bracket), *omit} try: @@ -1848,6 +1873,17 @@ def maybe_make_parens_invisible_in_atom( first = node.children[0] last = node.children[-1] + if ( + Preview.concatenated_list_comprehensions in mode + and is_lpar_token(first) + and is_rpar_token(last) + and first.value + and last.value + and len(node.children) >= 3 + and is_concatenated_list_comprehension(node.children[1]) + ): + return False + if is_lpar_token(first) and is_rpar_token(last): middle = node.children[1] # make parentheses invisible diff --git a/src/black/mode.py b/src/black/mode.py index 702f580e979..3ecc831770a 100644 --- a/src/black/mode.py +++ b/src/black/mode.py @@ -232,6 +232,7 @@ class Preview(Enum): fix_fmt_skip_in_one_liners = auto() standardize_type_comments = auto() wrap_comprehension_in = auto() + concatenated_list_comprehensions = 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/nodes.py b/src/black/nodes.py index 96bc20f20b3..a1b2afddfc0 100644 --- a/src/black/nodes.py +++ b/src/black/nodes.py @@ -645,6 +645,55 @@ def is_generator(node: LN) -> bool: return any(child.type == syms.old_comp_for for child in gexp.children) +def _listmaker_is_comprehension(node: LN) -> bool: + return isinstance(node, Node) and node.type == syms.listmaker and any( + isinstance(child, Node) and child.type in {syms.comp_for, syms.old_comp_for} + for child in node.children + ) + + +def is_list_comprehension(node: LN) -> bool: + """Return True if `node` represents a list comprehension expression.""" + if isinstance(node, Leaf): + return False + + if node.type == syms.atom and len(node.children) >= 3: + first = node.children[0] + last = node.children[-1] + if first.type == token.LSQB and last.type == token.RSQB: + return _listmaker_is_comprehension(node.children[1]) + if first.type == token.LPAR and last.type == token.RPAR: + inner = unwrap_singleton_parenthesis(node) + return inner is not None and is_list_comprehension(inner) + + if node.type == syms.power: + return any(is_list_comprehension(child) for child in node.children) + + return False + + +def is_concatenated_list_comprehension(node: LN) -> bool: + """Return True if `node` is an addition of list comprehensions.""" + if isinstance(node, Leaf): + return False + + if node.type == syms.atom: + inner = unwrap_singleton_parenthesis(node) + return inner is not None and is_concatenated_list_comprehension(inner) + + if node.type != syms.arith_expr: + return False + + operands = node.children[::2] + operators = node.children[1::2] + if len(operands) <= 1: + return False + + return all( + isinstance(op, Leaf) and op.type == token.PLUS for op in operators + ) and all(is_list_comprehension(operand) for operand in operands) + + def is_one_sequence_between( opening: Leaf, closing: Leaf, diff --git a/src/black/resources/black.schema.json b/src/black/resources/black.schema.json index bed70a4bb22..35d2d6a0ad6 100644 --- a/src/black/resources/black.schema.json +++ b/src/black/resources/black.schema.json @@ -88,6 +88,7 @@ "fix_fmt_skip_in_one_liners", "standardize_type_comments", "wrap_comprehension_in", + "concatenated_list_comprehensions", "remove_parens_around_except_types", "normalize_cr_newlines", "fix_module_docstring_detection", diff --git a/tests/data/cases/preview_concatenated_list_comprehensions.py b/tests/data/cases/preview_concatenated_list_comprehensions.py new file mode 100644 index 00000000000..f1a43303d33 --- /dev/null +++ b/tests/data/cases/preview_concatenated_list_comprehensions.py @@ -0,0 +1,27 @@ +# flags: --preview --line-length=120 + +matching_routes = [route for route in network_routes if route.destination and router_ip in route.destination_network] + [route for route in network_routes if route.destination is None and route.family == requested_family] + +already_wrapped = ( + [route for route in network_routes if route.load_balanced] + + [route for route in network_routes if route.fallback_route] +) + +triple = [route for route in network_routes if route.has_primary and route.supports_failover] + [route for route in network_routes if route.has_secondary and route.supports_failover] + [route for route in network_routes if route.is_backup and route.supports_failover] + +# output +matching_routes = ( + [route for route in network_routes if route.destination and router_ip in route.destination_network] + + [route for route in network_routes if route.destination is None and route.family == requested_family] +) + +already_wrapped = ( + [route for route in network_routes if route.load_balanced] + + [route for route in network_routes if route.fallback_route] +) + +triple = ( + [route for route in network_routes if route.has_primary and route.supports_failover] + + [route for route in network_routes if route.has_secondary and route.supports_failover] + + [route for route in network_routes if route.is_backup and route.supports_failover] +)