diff --git a/crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/named_expr.py b/crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/named_expr.py index 58522f5ad9..2b9d0ade31 100644 --- a/crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/named_expr.py +++ b/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) diff --git a/crates/ruff_python_formatter/src/comments/placement.rs b/crates/ruff_python_formatter/src/comments/placement.rs index df2a0a1cfc..c67647227b 100644 --- a/crates/ruff_python_formatter/src/comments/placement.rs +++ b/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( 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 8858c6079a..703b7d5d70 100644 --- a/crates/ruff_python_formatter/src/expression/expr_named_expr.rs +++ b/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 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(), - space(), - text(":="), - space(), - value.format(), + group(&format_args!(target.format(), soft_line_break_or_space())), + text(":=") ] - ) + )?; + + 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(()) } } 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 8a57c6cb90..95a899f8f2 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 @@ -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