Skip to content

Commit

Permalink
also handle end of line comments trailing a lone walrus
Browse files Browse the repository at this point in the history
  • Loading branch information
davidszotten committed Jul 12, 2023
1 parent 06fef9f commit 982b712
Show file tree
Hide file tree
Showing 3 changed files with 40 additions and 12 deletions.
36 changes: 27 additions & 9 deletions crates/ruff_python_formatter/src/comments/placement.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down Expand Up @@ -1233,36 +1233,34 @@ 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> {
let AnyNodeRef::ExprNamedExpr(named) = comment.enclosing_node() else {
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
// [
Expand All @@ -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)
Expand Down
13 changes: 12 additions & 1 deletion crates/ruff_python_formatter/src/expression/expr_named_expr.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
use crate::comments::dangling_comments;
use crate::context::PyFormatContext;
use crate::expression::parentheses::{
default_expression_needs_parentheses, NeedsParentheses, Parentheses, Parenthesize,
Expand All @@ -23,14 +24,15 @@ 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()])?;
} else {
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()])?;
Expand All @@ -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 {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -33,8 +33,7 @@ if (
# 1
x # 2
# 2.5
# 3
:=
:= # 3
# 3.5
y # 4
):
Expand Down

0 comments on commit 982b712

Please sign in to comment.