Handle comments on open parentheses in with statements (#6515)

## Summary

This PR adds handling for comments on open parentheses in parenthesized
context managers. For example, given:

```python
with (  # comment
    CtxManager1() as example1,
    CtxManager2() as example2,
    CtxManager3() as example3,
):
    ...
```

We want to preserve that formatting. (Black does the same.) On `main`,
we format as:

```python
with (
    # comment
    CtxManager1() as example1,
    CtxManager2() as example2,
    CtxManager3() as example3,
):
    ...
```

It's very similar to how `StmtImportFrom` is handled.

Note that this case _isn't_ covered by the "parenthesized comment"
proposal, since this is a common on the statement that would typically
be attached to the first `WithItem`, and the `WithItem` _itself_ can
have parenthesized comments, like:

```python
with (  # comment
    (
        CtxManager1()  # comment
    ) as example1,
    CtxManager2() as example2,
    CtxManager3() as example3,
):
    ...
```

## Test Plan

`cargo test`

Confirmed no change in similarity score.

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

@@ -116,3 +116,21 @@ with [
     dddddddddddddddddddddddddddddddd,
 ] as example1, aaaaaaaaaaaaaaaaaaaaaaaaaa * bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb * cccccccccccccccccccccccccccc + ddddddddddddddddd as example2, CtxManager222222222222222() as example2:
     ...
+
+# Comments on open parentheses
+with (  # comment
+    CtxManager1() as example1,
+    CtxManager2() as example2,
+    CtxManager3() as example3,
+):
+    ...
+
+
+with (  # outer comment
+    (  # inner comment
+        CtxManager1()
+    ) as example1,
+    CtxManager2() as example2,
+    CtxManager3() as example3,
+):
+    ...

crates/ruff_python_formatter/src/comments/placement.rs

@@ -73,6 +73,7 @@ pub(super) fn place_comment<'a>(
             handle_leading_class_with_decorators_comment(comment, class_def)
         }
         AnyNodeRef::StmtImportFrom(import_from) => handle_import_from_comment(comment, import_from),
+        AnyNodeRef::StmtWith(with_) => handle_with_comment(comment, with_),
         AnyNodeRef::ExprConstant(_) => {
             if let Some(AnyNodeRef::ExprFString(fstring)) = comment.enclosing_parent() {
                 CommentPlacement::dangling(fstring, comment)
@@ -1172,7 +1173,7 @@ fn handle_bracketed_end_of_line_comment<'a>(
     CommentPlacement::Default(comment)
 }
 
-/// Attach an enclosed end-of-line comment to a [`StmtImportFrom`].
+/// Attach an enclosed end-of-line comment to a [`ast::StmtImportFrom`].
 ///
 /// For example, given:
 /// ```python
@@ -1181,8 +1182,8 @@ fn handle_bracketed_end_of_line_comment<'a>(
 /// )
 /// ```
 ///
-/// The comment will be attached to the `StmtImportFrom` node as a dangling comment, to ensure
+/// The comment will be attached to the [`ast::StmtImportFrom`] node as a dangling comment, to
-/// that it remains on the same line as the `StmtImportFrom` itself.
+/// ensure that it remains on the same line as the [`ast::StmtImportFrom`] itself.
 fn handle_import_from_comment<'a>(
     comment: DecoratedComment<'a>,
     import_from: &'a ast::StmtImportFrom,
@@ -1217,6 +1218,35 @@ fn handle_import_from_comment<'a>(
     }
 }
 
+/// Attach an enclosed end-of-line comment to a [`ast::StmtWith`].
+///
+/// For example, given:
+/// ```python
+/// with ( # foo
+///     CtxManager1() as example1,
+///     CtxManager2() as example2,
+///     CtxManager3() as example3,
+/// ):
+///     ...
+/// ```
+///
+/// The comment will be attached to the [`ast::StmtWith`] node as a dangling comment, to ensure
+/// that it remains on the same line as the [`ast::StmtWith`] itself.
+fn handle_with_comment<'a>(
+    comment: DecoratedComment<'a>,
+    with_statement: &'a ast::StmtWith,
+) -> CommentPlacement<'a> {
+    if comment.line_position().is_end_of_line()
+        && with_statement.items.first().is_some_and(|with_item| {
+            with_statement.start() < comment.start() && comment.start() < with_item.start()
+        })
+    {
+        CommentPlacement::dangling(comment.enclosing_node(), comment)
+    } else {
+        CommentPlacement::Default(comment)
+    }
+}
+
 // Handle comments inside comprehensions, e.g.
 //
 // ```python

crates/ruff_python_formatter/src/statement/stmt_import_from.rs

@@ -4,7 +4,7 @@ use ruff_python_ast::node::AstNode;
 use ruff_python_ast::{Ranged, StmtImportFrom};
 
 use crate::builders::{parenthesize_if_expands, PyFormatterExtensions};
-use crate::comments::trailing_comments;
+use crate::expression::parentheses::parenthesized;
 use crate::{AsFormat, FormatNodeRule, PyFormatter};
 
 #[derive(Default)]
@@ -15,8 +15,8 @@ impl FormatNodeRule<StmtImportFrom> for FormatStmtImportFrom {
         let StmtImportFrom {
             module,
             names,
-            range: _,
             level,
+            range: _,
         } = item;
 
         let level_str = level
@@ -43,16 +43,29 @@ impl FormatNodeRule<StmtImportFrom> for FormatStmtImportFrom {
             }
         }
 
-        let comments = f.context().comments().clone();
-        let dangling_comments = comments.dangling_comments(item.as_any_node_ref());
-        write!(f, [trailing_comments(dangling_comments)])?;
-
         let names = format_with(|f| {
             f.join_comma_separated(item.end())
                 .entries(names.iter().map(|name| (name, name.format())))
                 .finish()
         });
-        parenthesize_if_expands(&names).fmt(f)
+
+        // A dangling comment on an import is a parenthesized comment, like:
+        // ```python
+        // from example import (  # comment
+        //     A,
+        //     B,
+        // )
+        // ```
+        let comments = f.context().comments().clone();
+        let parenthesized_comments = comments.dangling_comments(item.as_any_node_ref());
+
+        if parenthesized_comments.is_empty() {
+            parenthesize_if_expands(&names).fmt(f)
+        } else {
+            parenthesized("(", &names, ")")
+                .with_dangling_comments(parenthesized_comments)
+                .fmt(f)
+        }
     }
 
     fn fmt_dangling_comments(

crates/ruff_python_formatter/src/statement/stmt_with.rs

@@ -1,11 +1,12 @@
 use ruff_formatter::{format_args, write, FormatError};
+use ruff_python_ast::node::AstNode;
 use ruff_python_ast::{Ranged, StmtWith};
 use ruff_python_trivia::{SimpleTokenKind, SimpleTokenizer};
 use ruff_text_size::TextRange;
 
 use crate::comments::trailing_comments;
 use crate::expression::parentheses::{
-    in_parentheses_only_soft_line_break_or_space, optional_parentheses,
+    in_parentheses_only_soft_line_break_or_space, optional_parentheses, parenthesized,
 };
 use crate::prelude::*;
 use crate::FormatNodeRule;
@@ -15,9 +16,6 @@ pub struct FormatStmtWith;
 
 impl FormatNodeRule<StmtWith> for FormatStmtWith {
     fn fmt_fields(&self, item: &StmtWith, f: &mut PyFormatter) -> FormatResult<()> {
-        let comments = f.context().comments().clone();
-        let dangling_comments = comments.dangling_comments(item);
-
         write!(
             f,
             [
@@ -28,7 +26,39 @@ impl FormatNodeRule<StmtWith> for FormatStmtWith {
             ]
         )?;
 
-        if are_with_items_parenthesized(item, f.context())? {
+        // The `with` statement can have one dangling comment on the open parenthesis, like:
+        // ```python
+        // with (  # comment
+        //     CtxManager() as example
+        // ):
+        //     ...
+        // ```
+        //
+        // Any other dangling comments are trailing comments on the colon, like:
+        // ```python
+        // with CtxManager() as example:  # comment
+        //     ...
+        // ```
+        let comments = f.context().comments().clone();
+        let dangling_comments = comments.dangling_comments(item.as_any_node_ref());
+        let partition_point = dangling_comments.partition_point(|comment| {
+            item.items
+                .first()
+                .is_some_and(|with_item| with_item.start() > comment.start())
+        });
+        let (parenthesized_comments, colon_comments) = dangling_comments.split_at(partition_point);
+
+        if !parenthesized_comments.is_empty() {
+            let joined = format_with(|f: &mut PyFormatter| {
+                f.join_comma_separated(item.body.first().unwrap().start())
+                    .nodes(&item.items)
+                    .finish()
+            });
+
+            parenthesized("(", &joined, ")")
+                .with_dangling_comments(parenthesized_comments)
+                .fmt(f)?;
+        } else if are_with_items_parenthesized(item, f.context())? {
             optional_parentheses(&format_with(|f| {
                 let mut joiner = f.join_comma_separated(item.body.first().unwrap().start());
 
@@ -52,7 +82,7 @@ impl FormatNodeRule<StmtWith> for FormatStmtWith {
             f,
             [
                 text(":"),
-                trailing_comments(dangling_comments),
+                trailing_comments(colon_comments),
                 block_indent(&item.body.format())
             ]
         )

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

@@ -170,8 +170,7 @@ async def h():
     )
 
 
-with (
+with (  # d 1
-    # d 1
     x
 ):
     pass

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

@@ -88,7 +88,9 @@ from a import (  # comment
     bar,
 )
 
-from a import bar  # comment
+from a import (  # comment
+    bar
+)
 
 from a import (
     bar,

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

@@ -122,6 +122,24 @@ with [
     dddddddddddddddddddddddddddddddd,
 ] as example1, aaaaaaaaaaaaaaaaaaaaaaaaaa * bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb * cccccccccccccccccccccccccccc + ddddddddddddddddd as example2, CtxManager222222222222222() as example2:
     ...
+
+# Comments on open parentheses
+with (  # comment
+    CtxManager1() as example1,
+    CtxManager2() as example2,
+    CtxManager3() as example3,
+):
+    ...
+
+
+with (  # outer comment
+    (  # inner comment
+        CtxManager1()
+    ) as example1,
+    CtxManager2() as example2,
+    CtxManager3() as example3,
+):
+    ...
 ```
 
 ## Output
@@ -250,6 +268,25 @@ with [
     dddddddddddddddddddddddddddddddd,
 ] as example1, aaaaaaaaaaaaaaaaaaaaaaaaaa * bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb * cccccccccccccccccccccccccccc + ddddddddddddddddd as example2, CtxManager222222222222222() as example2:
     ...
+
+# Comments on open parentheses
+with (  # comment
+    CtxManager1() as example1,
+    CtxManager2() as example2,
+    CtxManager3() as example3,
+):
+    ...
+
+
+with (  # outer comment
+    (
+        # inner comment
+        CtxManager1()
+    ) as example1,
+    CtxManager2() as example2,
+    CtxManager3() as example3,
+):
+    ...
 ```