From 982b712a7d8c81c0b573c6e2cf70f3dddf9f469b Mon Sep 17 00:00:00 2001
From: David Szotten <davidszotten@gmail.com>
Date: Wed, 12 Jul 2023 09:45:37 +0100
Subject: [PATCH] also handle end of line comments trailing a lone walrus

---
 .../src/comments/placement.rs                 | 36 ++++++++++++++-----
 .../src/expression/expr_named_expr.rs         | 13 ++++++-
 .../format@expression__named_expr.py.snap     |  3 +-
 3 files changed, 40 insertions(+), 12 deletions(-)

diff --git a/crates/ruff_python_formatter/src/comments/placement.rs b/crates/ruff_python_formatter/src/comments/placement.rs
index 11d584edfcfa3..2faaec8adca99 100644
--- a/crates/ruff_python_formatter/src/comments/placement.rs
+++ b/crates/ruff_python_formatter/src/comments/placement.rs
@@ -40,7 +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_comment_before_walrus,
+        handle_walrus_comments,
     ];
     for handler in HANDLERS {
         comment = match handler(comment, locator) {
@@ -1233,19 +1233,19 @@ fn handle_trailing_expression_starred_star_end_of_line_comment<'a>(
     CommentPlacement::leading(starred.as_any_node_ref(), comment)
 }
 
-/// Assign comments between the target and `:=` as trailing the target instead of leading the value
+/// 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_comment_before_walrus<'a>(
+fn handle_walrus_comments<'a>(
     comment: DecoratedComment<'a>,
     locator: &Locator,
 ) -> CommentPlacement<'a> {
@@ -1253,16 +1253,14 @@ fn handle_comment_before_walrus<'a>(
         return CommentPlacement::Default(comment);
     };
 
-    if !comment.line_position().is_own_line() {
-        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
     // [
@@ -1273,7 +1271,27 @@ fn handle_comment_before_walrus<'a>(
     //  ]
     // ```
     if comment.slice().end() < walrus.start() {
-        return CommentPlacement::trailing((&*named.target).into(), comment);
+        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)
diff --git a/crates/ruff_python_formatter/src/expression/expr_named_expr.rs b/crates/ruff_python_formatter/src/expression/expr_named_expr.rs
index a3a391bb7a027..7e56850ffb479 100644
--- a/crates/ruff_python_formatter/src/expression/expr_named_expr.rs
+++ b/crates/ruff_python_formatter/src/expression/expr_named_expr.rs
@@ -1,3 +1,4 @@
+use crate::comments::dangling_comments;
 use crate::context::PyFormatContext;
 use crate::expression::parentheses::{
     default_expression_needs_parentheses, NeedsParentheses, Parentheses, Parenthesize,
@@ -23,6 +24,7 @@ impl FormatNodeRule<ExprNamedExpr> for FormatExprNamedExpr {
         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()])?;
@@ -30,7 +32,7 @@ impl FormatNodeRule<ExprNamedExpr> for FormatExprNamedExpr {
             write!(f, [soft_line_break()])?;
         }
 
-        write!(f, [text(":=")])?;
+        write!(f, [text(":="), dangling_comments(dangling_item_comments)])?;
 
         if leading_value_comments.is_empty() {
             write!(f, [space()])?;
@@ -39,6 +41,15 @@ impl FormatNodeRule<ExprNamedExpr> for FormatExprNamedExpr {
         }
         write!(f, [value.format()])
     }
+
+    fn fmt_dangling_comments(
+        &self,
+        _node: &ExprNamedExpr,
+        _f: &mut PyFormatter,
+    ) -> FormatResult<()> {
+        // Handled in `fmt_fields`
+        Ok(())
+    }
 }
 
 impl NeedsParentheses for ExprNamedExpr {
diff --git a/crates/ruff_python_formatter/tests/snapshots/format@expression__named_expr.py.snap b/crates/ruff_python_formatter/tests/snapshots/format@expression__named_expr.py.snap
index 2503ff62374b5..c327062c58f0c 100644
--- a/crates/ruff_python_formatter/tests/snapshots/format@expression__named_expr.py.snap
+++ b/crates/ruff_python_formatter/tests/snapshots/format@expression__named_expr.py.snap
@@ -33,8 +33,7 @@ if (
     # 1
     x  # 2
     # 2.5
-    # 3
-    :=
+    :=  # 3
     # 3.5
     y  # 4
 ):