[WIP] Implement the wrap_comprehension_in preview style from Black

crates/ruff_python_formatter/src/expression/mod.rs

@@ -359,12 +359,14 @@ impl Format<PyFormatContext<'_>> for MaybeParenthesizeExpression<'_> {
             parenthesize,
         } = self;
 
-        let preserve_parentheses = parenthesize.is_optional()
-            && is_expression_parenthesized(
-                (*expression).into(),
-                f.context().comments().ranges(),
-                f.context().source(),
-            );
+        let preserve_parentheses = matches!(
+            parenthesize,
+            Parenthesize::Optional | Parenthesize::IfBreaksComprehension
+        ) && is_expression_parenthesized(
+            (*expression).into(),
+            f.context().comments().ranges(),
+            f.context().source(),
+        );
 
         // If we want to preserve parentheses, short-circuit.
         if preserve_parentheses {
@@ -387,7 +389,9 @@ impl Format<PyFormatContext<'_>> for MaybeParenthesizeExpression<'_> {
             // Therefore, it is unnecessary to add an additional pair of parentheses if an outer expression
             // is parenthesized. Unless, it's the `Parenthesize::IfBreaksParenthesizedNested` layout
             // where parenthesizing nested `maybe_parenthesized_expression` is explicitly desired.
-            _ if f.context().node_level().is_parenthesized() => {
+            _ if f.context().node_level().is_parenthesized()
+                && !matches!(parenthesize, Parenthesize::IfBreaksComprehension) =>
+            {
                 return if matches!(parenthesize, Parenthesize::IfBreaksParenthesizedNested) {
                     parenthesize_if_expands(&expression.format().with_options(Parentheses::Never))
                         .with_indent(!is_expression_huggable(expression, f.context()))
@@ -408,7 +412,8 @@ impl Format<PyFormatContext<'_>> for MaybeParenthesizeExpression<'_> {
                 Parenthesize::Optional
                 | Parenthesize::IfBreaks
                 | Parenthesize::IfBreaksParenthesized
-                | Parenthesize::IfBreaksParenthesizedNested => {
+                | Parenthesize::IfBreaksParenthesizedNested
+                | Parenthesize::IfBreaksComprehension => {
                     if can_omit_optional_parentheses(expression, f.context()) {
                         optional_parentheses(&unparenthesized).fmt(f)
                     } else {
@@ -417,7 +422,9 @@ impl Format<PyFormatContext<'_>> for MaybeParenthesizeExpression<'_> {
                 }
             },
             OptionalParentheses::BestFit => match parenthesize {
-                Parenthesize::IfBreaksParenthesized | Parenthesize::IfBreaksParenthesizedNested => {
+                Parenthesize::IfBreaksParenthesized
+                | Parenthesize::IfBreaksParenthesizedNested
+                | Parenthesize::IfBreaksComprehension => {
                     // Can-omit layout is relevant for `"abcd".call`. We don't want to add unnecessary
                     // parentheses in this case.
                     if can_omit_optional_parentheses(expression, f.context()) {
@@ -448,7 +455,8 @@ impl Format<PyFormatContext<'_>> for MaybeParenthesizeExpression<'_> {
                 | Parenthesize::IfBreaks
                 | Parenthesize::IfRequired
                 | Parenthesize::IfBreaksParenthesized
-                | Parenthesize::IfBreaksParenthesizedNested => unparenthesized.fmt(f),
+                | Parenthesize::IfBreaksParenthesizedNested
+                | Parenthesize::IfBreaksComprehension => unparenthesized.fmt(f),
             },
 
             OptionalParentheses::Always => {
@@ -989,6 +997,10 @@ impl OwnParentheses {
     const fn is_non_empty(self) -> bool {
         matches!(self, OwnParentheses::NonEmpty)
     }
+
+    pub(crate) const fn is_empty(self) -> bool {
+        matches!(self, OwnParentheses::Empty)
+    }
 }
 
 /// Returns the [`OwnParentheses`] value for a given [`Expr`], to indicate whether it has its

crates/ruff_python_formatter/src/expression/parentheses.rs

@@ -69,12 +69,12 @@ pub(crate) enum Parenthesize {
     /// [`maybe_parenthesized_expression`] calls unlike other layouts that always omit parentheses
     /// when outer parentheses are present.
     IfBreaksParenthesizedNested,
-}
 
-impl Parenthesize {
-    pub(crate) const fn is_optional(self) -> bool {
-        matches!(self, Parenthesize::Optional)
-    }
+    /// Parenthesize the expression if it doesn't fit on a line or if the expression is
+    /// parenthesized in the source code. This is similar to `Self::Optional` except that this
+    /// variant will also include parentheses in nested calls, like
+    /// `Self::IfBreaksParenthesizedNested`.
+    IfBreaksComprehension,
 }
 
 /// Whether it is necessary to add parentheses around an expression.

crates/ruff_python_formatter/src/other/comprehension.rs

@@ -5,8 +5,10 @@ use ruff_text_size::{Ranged, TextRange};
 
 use crate::comments::{leading_comments, trailing_comments};
 use crate::expression::expr_tuple::TupleParentheses;
-use crate::expression::parentheses::is_expression_parenthesized;
+use crate::expression::parentheses::{Parenthesize, is_expression_parenthesized};
+use crate::expression::{OwnParentheses, has_own_parentheses, maybe_parenthesize_expression};
 use crate::prelude::*;
+use crate::preview::is_wrap_comprehension_in_enabled;
 
 #[derive(Default)]
 pub struct FormatComprehension;
@@ -104,14 +106,32 @@ impl FormatNodeRule<Comprehension> for FormatComprehension {
                 leading_comments(before_in_comments),
                 token("in"),
                 trailing_comments(trailing_in_comments),
-                Spacer {
-                    expression: iter,
-                    preserve_parentheses: true
-                },
-                iter.format(),
             ]
         )?;
 
+        if is_wrap_comprehension_in_enabled(f.context())
+            && has_own_parentheses(iter, f.context()).is_none_or(OwnParentheses::is_empty)
+        {
+            write!(
+                f,
+                [
+                    space(),
+                    maybe_parenthesize_expression(iter, item, Parenthesize::IfBreaksComprehension)
+                ]
+            )?;
+        } else {
+            write!(
+                f,
+                [
+                    Spacer {
+                        expression: iter,
+                        preserve_parentheses: true
+                    },
+                    iter.format(),
+                ]
+            )?;
+        }
+
         if !ifs.is_empty() {
             let joined = format_with(|f| {
                 let mut joiner = f.join_with(soft_line_break_or_space());

crates/ruff_python_formatter/src/preview.rs

@@ -36,3 +36,9 @@ pub(crate) const fn is_remove_parens_around_except_types_enabled(
 ) -> bool {
     context.is_preview()
 }
+
+/// Returns `true` if the [`wrap_comprehension_in`](https://github.com/astral-sh/ruff/pull/TODO)
+/// preview style is enabled.
+pub(crate) const fn is_wrap_comprehension_in_enabled(context: &PyFormatContext) -> bool {
+    context.is_preview()
+}

crates/ruff_python_formatter/tests/snapshots/black_compatibility@cases__preview_wrap_comprehension_in.py.snap

@@ -81,83 +81,18 @@ dict_with_really_long_names = {
 ```diff
 --- Black
 +++ Ruff
-@@ -1,20 +1,14 @@
+@@ -46,7 +46,9 @@
  [
-     a
--    for graph_path_expression in (
--        refined_constraint.condition_as_predicate.variables
--    )
-+    for graph_path_expression in refined_constraint.condition_as_predicate.variables
- ]
- [
-     a
--    for graph_path_expression in (
--        refined_constraint.condition_as_predicate.variables
--    )
-+    for graph_path_expression in refined_constraint.condition_as_predicate.variables
- ]
- [
-     a
--    for graph_path_expression in (
--        refined_constraint.condition_as_predicate.variables
--    )
-+    for graph_path_expression in refined_constraint.condition_as_predicate.variables
- ]
- [
-     a
-@@ -25,9 +19,7 @@
-
- [
-     (foobar_very_long_key, foobar_very_long_value)
--    for foobar_very_long_key, foobar_very_long_value in (
--        foobar_very_long_dictionary.items()
--    )
-+    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
-@@ -47,9 +39,7 @@
      [
          x
-         for x in bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb
--        for y in (
--            xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
--        )
-+        for y in xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
-     ]
- ]
-
-@@ -58,31 +48,23 @@
-     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
--    )
-+    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()
--    )
-+    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
--    )
-+    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
--    )
-+    for key_with_super_really_long_name in dictionary_with_super_really_long_name
+-        for x in bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb
++        for x in (
++            bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb
++        )
+         for y in (
+             xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
+         )
+@@ -84,5 +86,5 @@
  }
  {
      key_with_super_really_long_name: key_with_super_really_long_name
@@ -171,15 +106,21 @@ dict_with_really_long_names = {
 ```python
 [
     a
-    for graph_path_expression in refined_constraint.condition_as_predicate.variables
+    for graph_path_expression in (
+        refined_constraint.condition_as_predicate.variables
+    )
 ]
 [
     a
-    for graph_path_expression in refined_constraint.condition_as_predicate.variables
+    for graph_path_expression in (
+        refined_constraint.condition_as_predicate.variables
+    )
 ]
 [
     a
-    for graph_path_expression in refined_constraint.condition_as_predicate.variables
+    for graph_path_expression in (
+        refined_constraint.condition_as_predicate.variables
+    )
 ]
 [
     a
@@ -190,7 +131,9 @@ dict_with_really_long_names = {
 
 [
     (foobar_very_long_key, foobar_very_long_value)
-    for foobar_very_long_key, foobar_very_long_value in foobar_very_long_dictionary.items()
+    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
@@ -209,8 +152,12 @@ expected = [i for i in (a if b else c)]
 [
     [
         x
-        for x in bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb
-        for y in xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
+        for x in (
+            bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb
+        )
+        for y in (
+            xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
+        )
     ]
 ]
 
@@ -219,21 +166,29 @@ 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
+    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()
+    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
+    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
+    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