Fix formatting of single with-item with trailing comment (#14005)

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

@@ -328,3 +328,42 @@ with aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb
 if True:
     with anyio.CancelScope(shield=True) if get_running_loop() else contextlib.nullcontext() as b:
         pass
+
+
+# Regression test for https://github.com/astral-sh/ruff/issues/14001
+with (
+    open(
+        "some/path.txt",
+        "rb",
+    )
+    if True
+    else open("other/path.txt")
+    # Bar
+):
+    pass
+
+
+with (  # trailing comment
+    open(
+        "some/path.txt",
+        "rb",
+    )
+    if True
+    else open("other/path.txt")
+    # Bar
+):
+    pass
+
+
+with (
+    (
+        open(
+            "some/path.txt",
+            "rb",
+        )
+        if True
+        else open("other/path.txt")
+    )
+    # Bar
+):
+    pass

crates/ruff_python_formatter/src/other/with_item.rs

@@ -6,9 +6,11 @@ use crate::expression::parentheses::{
     is_expression_parenthesized, parenthesized, Parentheses, Parenthesize,
 };
 use crate::prelude::*;
-use crate::preview::is_with_single_item_pre_39_enabled;
+use crate::preview::{
+    is_with_single_item_pre_39_enabled, is_with_single_target_parentheses_enabled,
+};
 
-#[derive(Copy, Clone, Debug, Eq, PartialEq, Default)]
+#[derive(Copy, Clone, Debug, Eq, PartialEq)]
 pub enum WithItemLayout {
     /// A with item that is the `with`s only context manager and its context expression is parenthesized.
     ///
@@ -69,8 +71,7 @@ pub enum WithItemLayout {
     ///     a
     // ): ...
     /// ```
-    #[default]
+    ParenthesizedContextManagers { single: bool },
-    ParenthesizedContextManagers,
 }
 
 #[derive(Default)]
@@ -78,6 +79,12 @@ pub struct FormatWithItem {
     layout: WithItemLayout,
 }
 
+impl Default for WithItemLayout {
+    fn default() -> Self {
+        WithItemLayout::ParenthesizedContextManagers { single: false }
+    }
+}
+
 impl FormatRuleWithOptions<WithItem, PyFormatContext<'_>> for FormatWithItem {
     type Options = WithItemLayout;
 
@@ -97,6 +104,10 @@ impl FormatNodeRule<WithItem> for FormatWithItem {
         let comments = f.context().comments().clone();
         let trailing_as_comments = comments.dangling(item);
 
+        // WARNING: The `is_parenthesized` returns false-positives
+        // if the `with` has a single item without a target.
+        // E.g., it returns `true` for `with (a)` even though the parentheses
+        // belong to the with statement and not the expression but it can't determine that.
         let is_parenthesized = is_expression_parenthesized(
             context_expr.into(),
             f.context().comments().ranges(),
@@ -105,10 +116,17 @@ impl FormatNodeRule<WithItem> for FormatWithItem {
 
         match self.layout {
             // Remove the parentheses of the `with_items` if the with statement adds parentheses
-            WithItemLayout::ParenthesizedContextManagers => {
+            WithItemLayout::ParenthesizedContextManagers { single } => {
-                if is_parenthesized {
+                // ...except if the with item is parenthesized and it's not the only with item or it has a target.
-                    // ...except if the with item is parenthesized, then use this with item as a preferred breaking point
+                // Then use the context expression as a preferred breaking point.
-                    // or when it has comments, then parenthesize it to prevent comments from moving.
+                let prefer_breaking_context_expression =
+                    if is_with_single_target_parentheses_enabled(f.context()) {
+                        (optional_vars.is_some() || !single) && is_parenthesized
+                    } else {
+                        is_parenthesized
+                    };
+
+                if prefer_breaking_context_expression {
                     maybe_parenthesize_expression(
                         context_expr,
                         item,

crates/ruff_python_formatter/src/preview.rs

@@ -72,3 +72,10 @@ pub(crate) fn is_docstring_code_block_in_docstring_indent_enabled(
 pub(crate) fn is_join_implicit_concatenated_string_enabled(context: &PyFormatContext) -> bool {
     context.is_preview()
 }
+
+/// Returns `true` if the bugfix for single-with items with a trailing comment targeting Python 3.9 or newer is enabled.
+///
+/// See [#14001](https://github.com/astral-sh/ruff/issues/14001)
+pub(crate) fn is_with_single_target_parentheses_enabled(context: &PyFormatContext) -> bool {
+    context.is_preview()
+}

crates/ruff_python_formatter/src/statement/stmt_with.rs

@@ -71,7 +71,10 @@ impl FormatNodeRule<StmtWith> for FormatStmtWith {
 
                         match layout {
                             WithItemsLayout::SingleWithTarget(single) => {
-                                optional_parentheses(&single.format()).fmt(f)
+                                optional_parentheses(&single.format().with_options(
+                                    WithItemLayout::ParenthesizedContextManagers { single: true },
+                                ))
+                                .fmt(f)
                             }
 
                             WithItemsLayout::SingleWithoutTarget(single) => single
@@ -93,7 +96,11 @@ impl FormatNodeRule<StmtWith> for FormatStmtWith {
                                     for item in &with_stmt.items {
                                         joiner.entry_with_line_separator(
                                             item,
-                                            &item.format(),
+                                            &item.format().with_options(
+                                                WithItemLayout::ParenthesizedContextManagers {
+                                                    single: with_stmt.items.len() == 1,
+                                                },
+                                            ),
                                             soft_line_break_or_space(),
                                         );
                                     }
@@ -114,9 +121,22 @@ impl FormatNodeRule<StmtWith> for FormatStmtWith {
                             WithItemsLayout::Parenthesized => parenthesized(
                                 "(",
                                 &format_with(|f: &mut PyFormatter| {
-                                    f.join_comma_separated(with_stmt.body.first().unwrap().start())
+                                    let mut joiner = f.join_comma_separated(
-                                        .nodes(&with_stmt.items)
+                                        with_stmt.body.first().unwrap().start(),
-                                        .finish()
+                                    );
+
+                                    for item in &with_stmt.items {
+                                        joiner.entry(
+                                            item,
+                                            &item.format().with_options(
+                                                WithItemLayout::ParenthesizedContextManagers {
+                                                    single: with_stmt.items.len() == 1,
+                                                },
+                                            ),
+                                        );
+                                    }
+
+                                    joiner.finish()
                                 }),
                                 ")",
                             )

crates/ruff_python_formatter/tests/snapshots/format@statement__with.py.snap

@@ -1,6 +1,7 @@
 ---
 source: crates/ruff_python_formatter/tests/fixtures.rs
 input_file: crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/with.py
+snapshot_kind: text
 ---
 ## Input
 ```python
@@ -334,6 +335,45 @@ with aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb
 if True:
     with anyio.CancelScope(shield=True) if get_running_loop() else contextlib.nullcontext() as b:
         pass
+
+
+# Regression test for https://github.com/astral-sh/ruff/issues/14001
+with (
+    open(
+        "some/path.txt",
+        "rb",
+    )
+    if True
+    else open("other/path.txt")
+    # Bar
+):
+    pass
+
+
+with (  # trailing comment
+    open(
+        "some/path.txt",
+        "rb",
+    )
+    if True
+    else open("other/path.txt")
+    # Bar
+):
+    pass
+
+
+with (
+    (
+        open(
+            "some/path.txt",
+            "rb",
+        )
+        if True
+        else open("other/path.txt")
+    )
+    # Bar
+):
+    pass
 ```
 
 ## Outputs
@@ -710,6 +750,49 @@ if True:
         shield=True
     ) if get_running_loop() else contextlib.nullcontext() as b:
         pass
+
+
+# Regression test for https://github.com/astral-sh/ruff/issues/14001
+with (
+    (
+        open(
+            "some/path.txt",
+            "rb",
+        )
+        if True
+        else open("other/path.txt")
+    )
+    # Bar
+):
+    pass
+
+
+with (  # trailing comment
+    (
+        open(
+            "some/path.txt",
+            "rb",
+        )
+        if True
+        else open("other/path.txt")
+    )
+    # Bar
+):
+    pass
+
+
+with (
+    (
+        open(
+            "some/path.txt",
+            "rb",
+        )
+        if True
+        else open("other/path.txt")
+    )
+    # Bar
+):
+    pass
 ```
 
 
@@ -787,7 +870,7 @@ if True:
          pass
  
  with Child(aaaaaaaaa, bbbbbbbbbbbbbbb, cccccc), Document(
-@@ -344,14 +356,20 @@
+@@ -344,57 +356,57 @@
  ):
      pass
  
@@ -813,6 +896,64 @@ if True:
 +        else contextlib.nullcontext()
 +    ) as b:
          pass
+ 
+ 
+ # Regression test for https://github.com/astral-sh/ruff/issues/14001
+ with (
+-    (
+-        open(
+-            "some/path.txt",
+-            "rb",
+-        )
+-        if True
+-        else open("other/path.txt")
++    open(
++        "some/path.txt",
++        "rb",
+     )
++    if True
++    else open("other/path.txt")
+     # Bar
+ ):
+     pass
+ 
+ 
+ with (  # trailing comment
+-    (
+-        open(
+-            "some/path.txt",
+-            "rb",
+-        )
+-        if True
+-        else open("other/path.txt")
++    open(
++        "some/path.txt",
++        "rb",
+     )
++    if True
++    else open("other/path.txt")
+     # Bar
+ ):
+     pass
+ 
+ 
+ with (
+-    (
+-        open(
+-            "some/path.txt",
+-            "rb",
+-        )
+-        if True
+-        else open("other/path.txt")
++    open(
++        "some/path.txt",
++        "rb",
+     )
++    if True
++    else open("other/path.txt")
+     # Bar
+ ):
+     pass
 ```
 
 
@@ -1237,4 +1378,41 @@ if True:
         else contextlib.nullcontext() as b
     ):
         pass
+
+
+# Regression test for https://github.com/astral-sh/ruff/issues/14001
+with (
+    open(
+        "some/path.txt",
+        "rb",
+    )
+    if True
+    else open("other/path.txt")
+    # Bar
+):
+    pass
+
+
+with (  # trailing comment
+    open(
+        "some/path.txt",
+        "rb",
+    )
+    if True
+    else open("other/path.txt")
+    # Bar
+):
+    pass
+
+
+with (
+    open(
+        "some/path.txt",
+        "rb",
+    )
+    if True
+    else open("other/path.txt")
+    # Bar
+):
+    pass
 ```