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`
This commit is contained in:
Charlie Marsh 2023-08-17 23:10:45 -04:00 committed by GitHub
parent a128fe5148
commit 26bba11be6
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 194 additions and 9 deletions

View file

@ -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(

View file

@ -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(),
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(())
}
}