Refactor the comment handling of a statement's last expression (#8920)

crates/ruff_python_formatter/src/expression/mod.rs

@@ -12,9 +12,7 @@ use ruff_python_trivia::CommentRanges;
 use ruff_text_size::Ranged;
 
 use crate::builders::parenthesize_if_expands;
-use crate::comments::{
+use crate::comments::{leading_comments, trailing_comments, LeadingDanglingTrailingComments};
-    leading_comments, trailing_comments, LeadingDanglingTrailingComments, SourceComment,
-};
 use crate::context::{NodeLevel, WithNodeLevel};
 use crate::expression::expr_generator_exp::is_generator_parenthesized;
 use crate::expression::expr_tuple::is_tuple_parenthesized;
@@ -434,106 +432,16 @@ impl Format<PyFormatContext<'_>> for MaybeParenthesizeExpression<'_> {
                 }
 
                 Parenthesize::IfBreaks => {
-                    // Is the expression the last token in the parent statement.
+                    if node_comments.has_trailing() {
-                    // Excludes `await` and `yield` for which Black doesn't seem to apply the layout?
+                        expression.format().with_options(Parentheses::Always).fmt(f)
-                    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,
-                        )
                     } 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| {
+                        best_fit_parenthesize(&expression.format().with_options(Parentheses::Never))
-                            inline_comments.mark_formatted();
+                            .with_group_id(Some(group_id))
-
+                            .fmt(f)
-                            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)
                     }
                 }
             },
@@ -1248,41 +1156,3 @@ impl From<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)
-            ]
-        )
-    }
-}

crates/ruff_python_formatter/src/statement/stmt_ann_assign.rs

@@ -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)]
@@ -33,7 +31,7 @@ impl FormatNodeRule<StmtAnnAssign> for FormatStmtAnnAssign {
                     space(),
                     token("="),
                     space(),
-                    maybe_parenthesize_expression(value, item, Parenthesize::IfBreaks)
+                    FormatStatementsLastExpression::new(value, item)
                 ]
             )?;
         }

crates/ruff_python_formatter/src/statement/stmt_assign.rs

@@ -1,9 +1,11 @@
 use ruff_formatter::{format_args, write, FormatError};
-use ruff_python_ast::{Expr, StmtAssign};
+use ruff_python_ast::{AnyNodeRef, Expr, StmtAssign};
 
-use crate::comments::{SourceComment, SuppressionKind};
+use crate::comments::{trailing_comments, SourceComment, SuppressionKind};
 use crate::context::{NodeLevel, WithNodeLevel};
-use crate::expression::parentheses::{Parentheses, Parenthesize};
+use crate::expression::parentheses::{
+    NeedsParentheses, OptionalParentheses, Parentheses, Parenthesize,
+};
 use crate::expression::{has_own_parentheses, maybe_parenthesize_expression};
 use crate::prelude::*;
 use crate::statement::trailing_semicolon;
@@ -34,14 +36,7 @@ impl FormatNodeRule<StmtAssign> for FormatStmtAssign {
             ]
         )?;
 
-        write!(
+        FormatStatementsLastExpression::new(value, item).fmt(f)?;
-            f,
-            [maybe_parenthesize_expression(
-                value,
-                item,
-                Parenthesize::IfBreaks
-            )]
-        )?;
 
         if f.options().source_type().is_ipynb()
             && f.context().node_level().is_last_top_level_statement()
@@ -133,3 +128,188 @@ enum ParenthesizeTarget {
     Never,
     IfBreaks,
 }
+
+/// Formats the last expression in statements that start with a keyword (like `return`) or after an operator (assignments).
+///
+/// It avoids parenthesizing unsplittable values (like `None`, `True`, `False`, Names, a subset of strings) just to make
+/// the trailing comment fit and inlines a trailing comment if the value itself exceeds the configured line width:
+///
+/// The implementation formats the statement's 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 = loooooooooooooooooooooooooooong # with_comment
+/// b = (
+///     short # with_comment
+/// )
+/// ```
+///
+/// Which gets formatted to:
+///
+/// ```python
+/// # formatted
+/// a = (
+///     loooooooooooooooooooooooooooong # with comment
+/// )
+/// b = short # with comment
+/// ```
+///
+/// The long name gets parenthesized because it exceeds the configured line width and the trailing comma of the
+/// statement gets formatted inside (instead of outside) the parentheses.
+///
+/// The `short` name gets unparenthesized because it fits into the configured line length, regardless of whether
+/// the comment exceeds the line width or not.
+///
+/// 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.
+pub(super) struct FormatStatementsLastExpression<'a> {
+    expression: &'a Expr,
+    parent: AnyNodeRef<'a>,
+}
+
+impl<'a> FormatStatementsLastExpression<'a> {
+    pub(super) fn new<P: Into<AnyNodeRef<'a>>>(expression: &'a Expr, parent: P) -> Self {
+        Self {
+            expression,
+            parent: parent.into(),
+        }
+    }
+}
+
+impl Format<PyFormatContext<'_>> for FormatStatementsLastExpression<'_> {
+    fn fmt(&self, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
+        let can_inline_comment = match self.expression {
+            Expr::Name(_)
+            | Expr::NoneLiteral(_)
+            | Expr::NumberLiteral(_)
+            | Expr::BooleanLiteral(_) => true,
+            Expr::StringLiteral(string) => {
+                string.needs_parentheses(self.parent, f.context()) == OptionalParentheses::BestFit
+            }
+            Expr::BytesLiteral(bytes) => {
+                bytes.needs_parentheses(self.parent, f.context()) == OptionalParentheses::BestFit
+            }
+            Expr::FString(fstring) => {
+                fstring.needs_parentheses(self.parent, f.context()) == OptionalParentheses::BestFit
+            }
+            _ => false,
+        };
+
+        if !can_inline_comment {
+            return maybe_parenthesize_expression(
+                self.expression,
+                self.parent,
+                Parenthesize::IfBreaks,
+            )
+            .fmt(f);
+        }
+
+        let comments = f.context().comments().clone();
+        let expression_comments = comments.leading_dangling_trailing(self.expression);
+
+        if expression_comments.has_leading() {
+            // Preserve the parentheses if the expression has any leading comments,
+            // same as `maybe_parenthesize_expression`
+            return self
+                .expression
+                .format()
+                .with_options(Parentheses::Always)
+                .fmt(f);
+        }
+
+        let statement_trailing_comments = comments.trailing(self.parent);
+        let after_end_of_line = statement_trailing_comments
+            .partition_point(|comment| comment.line_position().is_end_of_line());
+        let (stmt_inline_comments, _) = statement_trailing_comments.split_at(after_end_of_line);
+
+        let after_end_of_line = expression_comments
+            .trailing
+            .partition_point(|comment| comment.line_position().is_end_of_line());
+
+        let (expression_inline_comments, expression_trailing_comments) =
+            expression_comments.trailing.split_at(after_end_of_line);
+
+        if expression_trailing_comments.is_empty() {
+            let inline_comments = OptionalParenthesesInlinedComments {
+                expression: expression_inline_comments,
+                statement: stmt_inline_comments,
+            };
+
+            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();
+
+                self.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 {
+            self.expression
+                .format()
+                .with_options(Parentheses::Always)
+                .fmt(f)
+        }
+    }
+}
+
+#[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.expression {
+            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)
+            ]
+        )
+    }
+}

crates/ruff_python_formatter/src/statement/stmt_aug_assign.rs

@@ -2,10 +2,8 @@ use ruff_formatter::write;
 use ruff_python_ast::StmtAugAssign;
 
 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;
 use crate::{AsFormat, FormatNodeRule};
 
@@ -28,7 +26,7 @@ impl FormatNodeRule<StmtAugAssign> for FormatStmtAugAssign {
                 op.format(),
                 token("="),
                 space(),
-                maybe_parenthesize_expression(value, item, Parenthesize::IfBreaks)
+                FormatStatementsLastExpression::new(value, item)
             ]
         )?;
 

crates/ruff_python_formatter/src/statement/stmt_return.rs

@@ -3,9 +3,8 @@ use ruff_python_ast::{Expr, StmtReturn};
 
 use crate::comments::{SourceComment, SuppressionKind};
 use crate::expression::expr_tuple::TupleParentheses;
-use crate::expression::maybe_parenthesize_expression;
-use crate::expression::parentheses::Parenthesize;
 use crate::prelude::*;
+use crate::statement::stmt_assign::FormatStatementsLastExpression;
 
 #[derive(Default)]
 pub struct FormatStmtReturn;
@@ -31,10 +30,7 @@ impl FormatNodeRule<StmtReturn> for FormatStmtReturn {
             Some(value) => {
                 write!(
                     f,
-                    [
+                    [space(), FormatStatementsLastExpression::new(value, item)]
-                        space(),
-                        maybe_parenthesize_expression(value, item, Parenthesize::IfBreaks)
-                    ]
                 )
             }
             None => Ok(()),