Parenthesize multi-context managers (#9222)

crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/with.options.json

@@ -0,0 +1,9 @@
+[
+  {
+    "target_version": "py38"
+  },
+  {
+    "target_version": "py39",
+    "preview": "enabled"
+  }
+]

crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/with_39.options.json

@@ -0,0 +1,6 @@
+[
+  {
+    "target_version": "py39",
+    "preview": "enabled"
+  }
+]

crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/with_39.py

@@ -0,0 +1,88 @@
+if True:
+    with (
+        anyio.CancelScope(shield=True)
+        if get_running_loop()
+        else contextlib.nullcontext()
+    ):
+        pass
+
+
+# Black avoids parenthesizing the with because it can make all with items fit by just breaking
+# around parentheses. We don't implement this optimisation because it makes it difficult to see where
+# the different context managers start and end.
+with cmd, xxxxxxxx.some_kind_of_method(
+    some_argument=[
+        "first",
+        "second",
+        "third",
+    ]
+) as cmd, another, and_more as x:
+    pass
+
+# Avoid parenthesizing single item context managers when splitting after the parentheses (can_omit_optional_parentheses)
+# is sufficient
+with xxxxxxxx.some_kind_of_method(
+    some_argument=[
+        "first",
+        "second",
+        "third",
+    ]
+).another_method(): pass
+
+if True:
+    with (
+        anyio.CancelScope(shield=True)
+        if get_running_loop()
+        else contextlib.nullcontext()
+    ):
+        pass
+
+
+# Black avoids parentheses here because it can make the entire with
+# header fit without requiring parentheses to do so.
+# We don't implement this optimisation because it very difficult to see where
+# the different context managers start or end.
+with cmd, xxxxxxxx.some_kind_of_method(
+    some_argument=[
+        "first",
+        "second",
+        "third",
+    ]
+) as cmd, another, and_more as x:
+    pass
+
+# Avoid parenthesizing single item context managers when splitting after the parentheses
+# is sufficient
+with xxxxxxxx.some_kind_of_method(
+    some_argument=[
+        "first",
+        "second",
+        "third",
+    ]
+).another_method(): pass
+
+# Parenthesize the with items if it makes them fit. Breaking the binary expression isn't
+# necessary because the entire items fit just into the 88 character limit.
+with aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb as c:
+    pass
+
+
+# Black parenthesizes this binary expression but also preserves the parentheses of the first with-item.
+# It does so because it prefers splitting already parenthesized context managers, even if it leads to more parentheses
+# like in this case.
+with (
+    (
+        aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb
+    ) as b,
+    c as d,
+):
+    pass
+
+if True:
+    with anyio.CancelScope(shield=True) if get_running_loop() else contextlib.nullcontext():
+        pass
+
+with (aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb as c):
+    pass
+
+

crates/ruff_python_formatter/src/expression/mod.rs

@@ -538,8 +538,8 @@ pub(crate) fn can_omit_optional_parentheses(expr: &Expr, context: &PyFormatConte
         false
     } else if visitor.max_precedence_count > 1 {
         false
-    } else if visitor.max_precedence == OperatorPrecedence::None && expr.is_lambda_expr() {
-        // Micha: This seems to exclusively apply for lambda expressions where the body ends in a subscript.
+    } else if visitor.max_precedence == OperatorPrecedence::None {
+        // Micha: This seems to apply for lambda expressions where the body ends in a subscript.
         // Subscripts are excluded by default because breaking them looks odd, but it seems to be fine for lambda expression.
         //
         // ```python
@@ -566,10 +566,19 @@ pub(crate) fn can_omit_optional_parentheses(expr: &Expr, context: &PyFormatConte
         //     ]
         // )
         // ```
+        //
+        // Another case are method chains:
+        // ```python
+        // xxxxxxxx.some_kind_of_method(
+        //     some_argument=[
+        //         "first",
+        //         "second",
+        //         "third",
+        //     ]
+        // ).another_method(a)
+        // ```
         true
-    } else if visitor.max_precedence == OperatorPrecedence::Attribute
-        && (expr.is_lambda_expr() || expr.is_named_expr_expr())
-    {
+    } else if visitor.max_precedence == OperatorPrecedence::Attribute {
         // A single method call inside a named expression (`:=`) or as the body of a lambda function:
         // ```python
         // kwargs["open_with"] = lambda path, _: fsspec.open(

crates/ruff_python_formatter/src/other/with_item.rs

@@ -1,5 +1,4 @@
 use ruff_formatter::write;
-
 use ruff_python_ast::WithItem;
 
 use crate::comments::SourceComment;
@@ -8,6 +7,7 @@ use crate::expression::parentheses::{
     is_expression_parenthesized, parenthesized, Parentheses, Parenthesize,
 };
 use crate::prelude::*;
+use crate::preview::is_wrap_multiple_context_managers_in_parens_enabled;
 
 #[derive(Default)]
 pub struct FormatWithItem;
@@ -23,26 +23,49 @@ impl FormatNodeRule<WithItem> for FormatWithItem {
         let comments = f.context().comments().clone();
         let trailing_as_comments = comments.dangling(item);
 
-        // Prefer keeping parentheses for already parenthesized expressions over
-        // parenthesizing other nodes.
-        let parenthesize = if is_expression_parenthesized(
+        let is_parenthesized = is_expression_parenthesized(
             context_expr.into(),
             f.context().comments().ranges(),
             f.context().source(),
-        ) {
-            Parenthesize::IfBreaks
+        );
+
+        // Remove the parentheses of the `with_items` if the with statement adds parentheses
+        if f.context().node_level().is_parenthesized()
+            && is_wrap_multiple_context_managers_in_parens_enabled(f.context())
+        {
+            if is_parenthesized {
+                // ...except if the with item is parenthesized, then use this with item as a preferred breaking point
+                // or when it has comments, then parenthesize it to prevent comments from moving.
+                maybe_parenthesize_expression(
+                    context_expr,
+                    item,
+                    Parenthesize::IfBreaksOrIfRequired,
+                )
+                .fmt(f)?;
+            } else {
+                context_expr
+                    .format()
+                    .with_options(Parentheses::Never)
+                    .fmt(f)?;
+            }
         } else {
-            Parenthesize::IfRequired
-        };
+            // Prefer keeping parentheses for already parenthesized expressions over
+            // parenthesizing other nodes.
+            let parenthesize = if is_parenthesized {
+                Parenthesize::IfBreaks
+            } else {
+                Parenthesize::IfRequired
+            };
 
-        write!(
-            f,
-            [maybe_parenthesize_expression(
-                context_expr,
-                item,
-                parenthesize
-            )]
-        )?;
+            write!(
+                f,
+                [maybe_parenthesize_expression(
+                    context_expr,
+                    item,
+                    parenthesize
+                )]
+            )?;
+        }
 
         if let Some(optional_vars) = optional_vars {
             write!(f, [space(), token("as"), space()])?;

crates/ruff_python_formatter/src/preview.rs

@@ -39,6 +39,15 @@ pub(crate) const fn is_no_blank_line_before_class_docstring_enabled(
     context.is_preview()
 }
 
+/// Returns `true` if the [`wrap_multiple_context_managers_in_parens`](https://github.com/astral-sh/ruff/issues/8889) preview style is enabled.
+///
+/// Unlike Black, we re-use the same preview style feature flag for [`improved_async_statements_handling`](https://github.com/astral-sh/ruff/issues/8890)
+pub(crate) const fn is_wrap_multiple_context_managers_in_parens_enabled(
+    context: &PyFormatContext,
+) -> bool {
+    context.is_preview()
+}
+
 /// Returns `true` if the [`module_docstring_newlines`](https://github.com/astral-sh/ruff/issues/7995) preview style is enabled.
 pub(crate) const fn is_module_docstring_newlines_enabled(context: &PyFormatContext) -> bool {
     context.is_preview()