Skip to content

Commit

Permalink
Refactor the comment handling of a statement's last expression
Browse files Browse the repository at this point in the history
  • Loading branch information
MichaReiser committed Nov 30, 2023
1 parent e8da95d commit 76110f7
Show file tree
Hide file tree
Showing 5 changed files with 203 additions and 161 deletions.
142 changes: 6 additions & 136 deletions crates/ruff_python_formatter/src/expression/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,9 +12,7 @@ use ruff_python_trivia::CommentRanges;
use ruff_text_size::Ranged;

use crate::builders::parenthesize_if_expands;
use crate::comments::{
leading_comments, trailing_comments, LeadingDanglingTrailingComments, SourceComment,
};
use crate::comments::{leading_comments, trailing_comments, LeadingDanglingTrailingComments};
use crate::context::{NodeLevel, WithNodeLevel};
use crate::expression::expr_generator_exp::is_generator_parenthesized;
use crate::expression::expr_tuple::is_tuple_parenthesized;
Expand Down Expand Up @@ -428,106 +426,16 @@ impl Format<PyFormatContext<'_>> for MaybeParenthesizeExpression<'_> {
expression.format().with_options(Parentheses::Never).fmt(f)
}
Parenthesize::IfBreaks => {
// Is the expression the last token in the parent statement.
// Excludes `await` and `yield` for which Black doesn't seem to apply the layout?
let last_expression = parent.is_stmt_assign()
|| parent.is_stmt_ann_assign()
|| parent.is_stmt_aug_assign()
|| parent.is_stmt_return();

// Format the statements and value's trailing end of line comments:
// * after the expression if the expression needs no parentheses (necessary or the `expand_parent` makes the group never fit).
// * inside the parentheses if the expression exceeds the line-width.
//
// ```python
// a = long # with_comment
// b = (
// short # with_comment
// )
//
// # formatted
// a = (
// long # with comment
// )
// b = short # with comment
// ```
// This matches Black's formatting with the exception that ruff applies this style also for
// attribute chains and non-fluent call expressions. See https://github.com/psf/black/issues/4001#issuecomment-1786681792
//
// This logic isn't implemented in [`place_comment`] by associating trailing statement comments to the expression because
// doing so breaks the suite empty lines formatting that relies on trailing comments to be stored on the statement.
let (inline_comments, expression_trailing_comments) = if last_expression
&& !(
// Ignore non-fluent attribute chains for black compatibility.
// See https://github.com/psf/black/issues/4001#issuecomment-1786681792
expression.is_attribute_expr()
|| expression.is_call_expr()
|| expression.is_yield_from_expr()
|| expression.is_yield_expr()
|| expression.is_await_expr()
) {
let parent_trailing_comments = comments.trailing(*parent);
let after_end_of_line = parent_trailing_comments
.partition_point(|comment| comment.line_position().is_end_of_line());
let (stmt_inline_comments, _) =
parent_trailing_comments.split_at(after_end_of_line);

let after_end_of_line = node_comments
.trailing
.partition_point(|comment| comment.line_position().is_end_of_line());

let (expression_inline_comments, expression_trailing_comments) =
node_comments.trailing.split_at(after_end_of_line);

(
OptionalParenthesesInlinedComments {
expression: expression_inline_comments,
statement: stmt_inline_comments,
},
expression_trailing_comments,
)
if node_comments.has_trailing() {
expression.format().with_options(Parentheses::Always).fmt(f)
} else {
(
OptionalParenthesesInlinedComments::default(),
node_comments.trailing,
)
};

if expression_trailing_comments.is_empty() {
// The group id is necessary because the nested expressions may reference it.
let group_id = f.group_id("optional_parentheses");
let f = &mut WithNodeLevel::new(NodeLevel::Expression(Some(group_id)), f);

best_fit_parenthesize(&format_with(|f| {
inline_comments.mark_formatted();

expression
.format()
.with_options(Parentheses::Never)
.fmt(f)?;

if !inline_comments.is_empty() {
// If the expressions exceeds the line width, format the comments in the parentheses
if_group_breaks(&inline_comments)
.with_group_id(Some(group_id))
.fmt(f)?;
}

Ok(())
}))
.with_group_id(Some(group_id))
.fmt(f)?;

if !inline_comments.is_empty() {
// If the line fits into the line width, format the comments after the parenthesized expression
if_group_fits_on_line(&inline_comments)
.with_group_id(Some(group_id))
.fmt(f)?;
}

Ok(())
} else {
expression.format().with_options(Parentheses::Always).fmt(f)
best_fit_parenthesize(&expression.format().with_options(Parentheses::Never))
.with_group_id(Some(group_id))
.fmt(f)
}
}
},
Expand Down Expand Up @@ -1161,41 +1069,3 @@ impl From<ast::Operator> for OperatorPrecedence {
}
}
}

#[derive(Debug, Default)]
struct OptionalParenthesesInlinedComments<'a> {
expression: &'a [SourceComment],
statement: &'a [SourceComment],
}

impl<'a> OptionalParenthesesInlinedComments<'a> {
fn is_empty(&self) -> bool {
self.expression.is_empty() && self.statement.is_empty()
}

fn iter_comments(&self) -> impl Iterator<Item = &'a SourceComment> {
self.expression.iter().chain(self.statement)
}

fn mark_formatted(&self) {
for comment in self.iter_comments() {
comment.mark_formatted();
}
}
}

impl Format<PyFormatContext<'_>> for OptionalParenthesesInlinedComments<'_> {
fn fmt(&self, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
for comment in self.iter_comments() {
comment.mark_unformatted();
}

write!(
f,
[
trailing_comments(self.expression),
trailing_comments(self.statement)
]
)
}
}
6 changes: 2 additions & 4 deletions crates/ruff_python_formatter/src/statement/stmt_ann_assign.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,8 @@ use ruff_formatter::write;
use ruff_python_ast::StmtAnnAssign;

use crate::comments::{SourceComment, SuppressionKind};

use crate::expression::maybe_parenthesize_expression;
use crate::expression::parentheses::Parenthesize;
use crate::prelude::*;
use crate::statement::stmt_assign::FormatStatementsLastExpression;
use crate::statement::trailing_semicolon;

#[derive(Default)]
Expand Down Expand Up @@ -33,7 +31,7 @@ impl FormatNodeRule<StmtAnnAssign> for FormatStmtAnnAssign {
space(),
token("="),
space(),
maybe_parenthesize_expression(value, item, Parenthesize::IfBreaks)
FormatStatementsLastExpression::new(value, item)
]
)?;
}
Expand Down
Loading

0 comments on commit 76110f7

Please sign in to comment.