Manually format comments around := in named expressions (#6634)

## Summary

Attaches comments around the `:=` operator in a named expression as
dangling, and formats them manually in the `named_expr.rs` formatter.

Closes https://github.com/astral-sh/ruff/issues/5695.

## Test Plan

`cargo test`

crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/named_expr.py

@@ -8,6 +8,45 @@ if (
 ):
     pass
 
+if (
+    # 1
+    x # 2
+    := # 3
+    (y) # 4
+):
+    pass
+
+if (
+    # 1
+    x # 2
+    := # 3
+    (y) # 4
+    # 5
+):
+    pass
+
+if (
+    # 1
+    x # 2
+    # 2.5
+    := # 3
+    # 3.5
+    y # 4
+):
+    pass
+
+if (
+    # 1
+    x # 2
+    # 2.5
+    := # 3
+    # 3.5
+    (  # 4
+        y # 5
+    )  # 6
+):
+    pass
+
 y0 = (y1 := f(x))
 
 f(x:=y, z=True)

crates/ruff_python_formatter/src/comments/placement.rs

@@ -181,6 +181,7 @@ fn handle_enclosed_comment<'a>(
             )
         }
         AnyNodeRef::Keyword(_) => handle_dict_unpacking_comment(comment, locator),
+        AnyNodeRef::ExprNamedExpr(_) => handle_named_expr_comment(comment, locator),
         AnyNodeRef::ExprDict(_) => handle_dict_unpacking_comment(comment, locator)
             .or_else(|comment| handle_bracketed_end_of_line_comment(comment, locator)),
         AnyNodeRef::ExprIfExp(expr_if) => handle_expr_if_comment(comment, expr_if, locator),
@@ -1105,7 +1106,7 @@ fn handle_trailing_expression_starred_star_end_of_line_comment<'a>(
 ///     # trailing a own line comment
 ///     as # trailing as same line comment
 ///     b
-// ): ...
+/// ): ...
 /// ```
 fn handle_with_item_comment<'a>(
     comment: DecoratedComment<'a>,
@@ -1138,6 +1139,49 @@ fn handle_with_item_comment<'a>(
     }
 }
 
+/// Handles comments around the `:=` token in a named expression (walrus operator).
+///
+/// For example, here, `# 1` and `# 2` will be marked as dangling comments on the named expression,
+/// while `# 3` and `4` will be attached `y` (via our general parenthesized comment handling), and
+/// `# 5` will be a trailing comment on the named expression.
+///
+/// ```python
+/// if (
+///     x
+///     :=  # 1
+///     # 2
+///     (  # 3
+///         y  # 4
+///     ) # 5
+/// ):
+///     pass
+/// ```
+fn handle_named_expr_comment<'a>(
+    comment: DecoratedComment<'a>,
+    locator: &Locator,
+) -> CommentPlacement<'a> {
+    debug_assert!(comment.enclosing_node().is_expr_named_expr());
+
+    let (Some(target), Some(value)) = (comment.preceding_node(), comment.following_node()) else {
+        return CommentPlacement::Default(comment);
+    };
+
+    let colon_equal = find_only_token_in_range(
+        TextRange::new(target.end(), value.start()),
+        SimpleTokenKind::ColonEqual,
+        locator,
+    );
+
+    if comment.end() < colon_equal.start() {
+        // If the comment is before the `:=` token, then it must be a trailing comment of the
+        // target.
+        CommentPlacement::trailing(target, comment)
+    } else {
+        // Otherwise, treat it as dangling. We effectively treat it as a comment on the `:=` itself.
+        CommentPlacement::dangling(comment.enclosing_node(), comment)
+    }
+}
+
 /// Looks for a token in the range that contains no other tokens except for parentheses outside
 /// the expression ranges
 fn find_only_token_in_range(

crates/ruff_python_formatter/src/expression/expr_named_expr.rs

@@ -1,7 +1,10 @@
+use crate::comments::{dangling_comments, SourceComment};
 use crate::context::PyFormatContext;
 use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses};
 use crate::{AsFormat, FormatNodeRule, PyFormatter};
-use ruff_formatter::prelude::{space, text};
+use ruff_formatter::prelude::{
+    format_args, group, hard_line_break, soft_line_break_or_space, space, text,
+};
 use ruff_formatter::{write, Buffer, FormatResult};
 use ruff_python_ast::node::AnyNodeRef;
 use ruff_python_ast::ExprNamedExpr;
@@ -16,16 +19,35 @@ impl FormatNodeRule<ExprNamedExpr> for FormatExprNamedExpr {
             value,
             range: _,
         } = item;
+
+        // This context, a dangling comment is an end-of-line comment on the same line as the `:=`.
+        let comments = f.context().comments().clone();
+        let dangling = comments.dangling_comments(item);
+
         write!(
             f,
             [
-                target.format(),
+                group(&format_args!(target.format(), soft_line_break_or_space())),
-                space(),
+                text(":=")
-                text(":="),
-                space(),
-                value.format(),
             ]
-        )
+        )?;
+
+        if dangling.is_empty() {
+            write!(f, [space()])?;
+        } else {
+            write!(f, [dangling_comments(dangling), hard_line_break()])?;
+        }
+
+        write!(f, [value.format()])
+    }
+
+    fn fmt_dangling_comments(
+        &self,
+        _dangling_comments: &[SourceComment],
+        _f: &mut PyFormatter,
+    ) -> FormatResult<()> {
+        // Handled by `fmt_fields`
+        Ok(())
     }
 }
 

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

@@ -14,6 +14,45 @@ if (
 ):
     pass
 
+if (
+    # 1
+    x # 2
+    := # 3
+    (y) # 4
+):
+    pass
+
+if (
+    # 1
+    x # 2
+    := # 3
+    (y) # 4
+    # 5
+):
+    pass
+
+if (
+    # 1
+    x # 2
+    # 2.5
+    := # 3
+    # 3.5
+    y # 4
+):
+    pass
+
+if (
+    # 1
+    x # 2
+    # 2.5
+    := # 3
+    # 3.5
+    (  # 4
+        y # 5
+    )  # 6
+):
+    pass
+
 y0 = (y1 := f(x))
 
 f(x:=y, z=True)
@@ -65,7 +104,48 @@ y = 1
 
 if (
     # 1
-    x := y  # 2  # 3  # 4
+    x  # 2
+    :=  # 3
+    y  # 4
+):
+    pass
+
+if (
+    # 1
+    x  # 2
+    :=  # 3
+    (y)  # 4
+):
+    pass
+
+if (
+    # 1
+    x  # 2
+    :=  # 3
+    (y)  # 4
+    # 5
+):
+    pass
+
+if (
+    # 1
+    x  # 2
+    # 2.5
+    :=  # 3
+    # 3.5
+    y  # 4
+):
+    pass
+
+if (
+    # 1
+    x  # 2
+    # 2.5
+    :=  # 3
+    # 3.5
+    (  # 4
+        y  # 5
+    )  # 6
 ):
     pass