Fix format named expr comments

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

@@ -3,10 +3,16 @@
 if (
     # 1
     x # 2
+    # 2.5
     := # 3
+    # 3.5
     y # 4
 ):
-    pass
+    ...
+
+# should not add brackets
+if x := y:
+    ...
 
 y0 = (y1 := f(x))
 

crates/ruff_python_formatter/src/comments/placement.rs

@@ -40,6 +40,7 @@ pub(super) fn place_comment<'a>(
         handle_expr_if_comment,
         handle_comprehension_comment,
         handle_trailing_expression_starred_star_end_of_line_comment,
+        handle_walrus_comments,
     ];
     for handler in HANDLERS {
         comment = match handler(comment, locator) {
@@ -1232,6 +1233,70 @@ fn handle_trailing_expression_starred_star_end_of_line_comment<'a>(
     CommentPlacement::leading(starred.as_any_node_ref(), comment)
 }
 
+/// Handle comments around `:=`
+///
+/// ```python
+/// if (
+///     x
+///     # make trailing x (instead of leading y)
+///     :=  # make dangling on the NamedExpr node
+///     y
+/// )
+///    ...
+/// ```
+// TODO: can this also handle comments before colon in dicts
+fn handle_walrus_comments<'a>(
+    comment: DecoratedComment<'a>,
+    locator: &Locator,
+) -> CommentPlacement<'a> {
+    let AnyNodeRef::ExprNamedExpr(named) = comment.enclosing_node() else {
+        return CommentPlacement::Default(comment);
+    };
+
+    let walrus = find_only_token_in_range(
+        TextRange::new(named.target.range().end(), named.value.range().start()),
+        locator,
+        TokenKind::Walrus,
+    );
+
+    let is_own_line = comment.line_position().is_own_line();
+
+    // Comments between the target and the `:=`
+    // ```python
+    // [
+    //      a
+    //      # attach as trailing on the target `a` (instead of leading on the value `b`)
+    //      :=
+    //      in b
+    //  ]
+    // ```
+    if comment.slice().end() < walrus.start() {
+        return if is_own_line {
+            CommentPlacement::trailing((&*named.target).into(), comment)
+        } else {
+            CommentPlacement::Default(comment)
+        };
+    }
+
+    // Comments after the `:=`
+    // ```python
+    // [
+    //      a
+    //      :=  # attach as dangling on the NamedExpr
+    //      in b
+    //  ]
+    // ```
+    if comment.slice().end() < named.value.range().start() {
+        return if is_own_line {
+            CommentPlacement::Default(comment)
+        } else {
+            return CommentPlacement::dangling(comment.enclosing_node(), comment);
+        };
+    }
+
+    CommentPlacement::Default(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(range: TextRange, locator: &Locator, token_kind: TokenKind) -> Token {

crates/ruff_python_formatter/src/expression/expr_named_expr.rs

@@ -1,9 +1,10 @@
+use crate::comments::dangling_comments;
 use crate::context::PyFormatContext;
 use crate::expression::parentheses::{
     default_expression_needs_parentheses, NeedsParentheses, Parentheses, Parenthesize,
 };
 use crate::{AsFormat, FormatNodeRule, PyFormatter};
-use ruff_formatter::prelude::{space, text};
+use ruff_formatter::prelude::{soft_line_break, space, text};
 use ruff_formatter::{write, Buffer, FormatResult};
 use rustpython_parser::ast::ExprNamedExpr;
 
@@ -17,16 +18,37 @@ impl FormatNodeRule<ExprNamedExpr> for FormatExprNamedExpr {
             value,
             range: _,
         } = item;
-        write!(
-            f,
-            [
-                target.format(),
-                space(),
-                text(":="),
-                space(),
-                value.format(),
-            ]
-        )
+
+        write!(f, [target.format()])?;
+
+        let comments = f.context().comments().clone();
+        let trailing_target_comments = comments.trailing_comments(target.as_ref());
+        let leading_value_comments = comments.leading_comments(value.as_ref());
+        let dangling_item_comments = comments.dangling_comments(item);
+
+        if trailing_target_comments.is_empty() {
+            write!(f, [space()])?;
+        } else {
+            write!(f, [soft_line_break()])?;
+        }
+
+        write!(f, [text(":="), dangling_comments(dangling_item_comments)])?;
+
+        if leading_value_comments.is_empty() {
+            write!(f, [space()])?;
+        } else {
+            write!(f, [soft_line_break()])?;
+        }
+        write!(f, [value.format()])
+    }
+
+    fn fmt_dangling_comments(
+        &self,
+        _node: &ExprNamedExpr,
+        _f: &mut PyFormatter,
+    ) -> FormatResult<()> {
+        // Handled in `fmt_fields`
+        Ok(())
     }
 }
 

crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__trivia__tests__tokenize_colon.snap

@@ -0,0 +1,18 @@
+---
+source: crates/ruff_python_formatter/src/trivia.rs
+expression: test_case.tokens()
+---
+[
+    Token {
+        kind: Colon,
+        range: 0..1,
+    },
+    Token {
+        kind: Colon,
+        range: 1..2,
+    },
+    Token {
+        kind: Colon,
+        range: 2..3,
+    },
+]

crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__trivia__tests__tokenize_walrus.snap

@@ -0,0 +1,14 @@
+---
+source: crates/ruff_python_formatter/src/trivia.rs
+expression: test_case.tokens()
+---
+[
+    Token {
+        kind: Walrus,
+        range: 0..2,
+    },
+    Token {
+        kind: Walrus,
+        range: 2..4,
+    },
+]

crates/ruff_python_formatter/src/trivia.rs

@@ -185,6 +185,10 @@ pub(crate) enum TokenKind {
     /// `.`.
     Dot,
 
+    /// `:=`.
+    /// TODO: name?
+    Walrus,
+
     /// `else`
     Else,
 
@@ -322,6 +326,15 @@ impl<'a> SimpleTokenizer<'a> {
 
             '\\' => TokenKind::Continuation,
 
+            ':' => {
+                if self.cursor.first() == '=' {
+                    self.cursor.bump();
+                    TokenKind::Walrus
+                } else {
+                    TokenKind::Colon
+                }
+            }
+
             c => {
                 let kind = if is_identifier_start(c) {
                     self.cursor.eat_while(is_identifier_continuation);
@@ -457,6 +470,9 @@ impl<'a> SimpleTokenizer<'a> {
                             self.cursor = savepoint;
                             TokenKind::Other
                         }
+                    } else if c == '=' && self.cursor.last() == ':' {
+                        self.cursor.bump_back();
+                        TokenKind::Walrus
                     } else {
                         TokenKind::from_non_trivia_char(c)
                     };
@@ -687,6 +703,26 @@ mod tests {
         test_case.assert_reverse_tokenization();
     }
 
+    #[test]
+    fn tokenize_colon() {
+        let source = ":::";
+
+        let test_case = tokenize(source);
+
+        assert_debug_snapshot!(test_case.tokens());
+        test_case.assert_reverse_tokenization();
+    }
+
+    #[test]
+    fn tokenize_walrus() {
+        let source = ":=:=";
+
+        let test_case = tokenize(source);
+
+        assert_debug_snapshot!(test_case.tokens());
+        test_case.assert_reverse_tokenization();
+    }
+
     #[test]
     fn tokenize_continuation() {
         let source = "( \\\n )";

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

@@ -9,10 +9,16 @@ y = 1
 if (
     # 1
     x # 2
+    # 2.5
     := # 3
+    # 3.5
     y # 4
 ):
-    pass
+    ...
+
+# should not add brackets
+if x := y:
+    ...
 
 y0 = (y1 := f(x))
 
@@ -25,9 +31,17 @@ y = 1
 
 if (
     # 1
-    x := y  # 2  # 3  # 4
+    x  # 2
+    # 2.5
+    :=  # 3
+    # 3.5
+    y  # 4
 ):
-    pass
+    ...
+
+# should not add brackets
+if (x := y):
+    ...
 
 y0 = (y1 := f(x))