Refactor needs parentheses

Author: MichaReiser

crates/ruff_python_formatter/src/builders.rs

@@ -1,9 +1,12 @@
+use ruff_text_size::{TextRange, TextSize};
+use rustpython_parser::ast::Ranged;
+
+use ruff_formatter::{format_args, write, Argument, Arguments};
+
 use crate::context::NodeLevel;
 use crate::prelude::*;
-use crate::trivia::{first_non_trivia_token, lines_after, skip_trailing_trivia, Token, TokenKind};
-use ruff_formatter::{format_args, write, Argument, Arguments};
-use ruff_text_size::TextSize;
-use rustpython_parser::ast::Ranged;
+use crate::trivia::{lines_after, skip_trailing_trivia, SimpleTokenizer, Token, TokenKind};
+use crate::MagicTrailingComma;
 
 /// Adds parentheses and indents `content` if it doesn't fit on a line.
 pub(crate) fn parenthesize_if_expands<'ast, T>(content: &T) -> ParenthesizeIfExpands<'_, 'ast>
@@ -53,16 +56,22 @@ pub(crate) trait PyFormatterExtensions<'ast, 'buf> {
     /// A builder that separates each element by a `,` and a [`soft_line_break_or_space`].
     /// It emits a trailing `,` that is only shown if the enclosing group expands. It forces the enclosing
     /// group to expand if the last item has a trailing `comma` and the magical comma option is enabled.
-    fn join_comma_separated<'fmt>(&'fmt mut self) -> JoinCommaSeparatedBuilder<'fmt, 'ast, 'buf>;
+    fn join_comma_separated<'fmt>(
+        &'fmt mut self,
+        sequence_end: TextSize,
+    ) -> JoinCommaSeparatedBuilder<'fmt, 'ast, 'buf>;
 }
 
 impl<'buf, 'ast> PyFormatterExtensions<'ast, 'buf> for PyFormatter<'ast, 'buf> {
     fn join_nodes<'fmt>(&'fmt mut self, level: NodeLevel) -> JoinNodesBuilder<'fmt, 'ast, 'buf> {
         JoinNodesBuilder::new(self, level)
     }
 
-    fn join_comma_separated<'fmt>(&'fmt mut self) -> JoinCommaSeparatedBuilder<'fmt, 'ast, 'buf> {
-        JoinCommaSeparatedBuilder::new(self)
+    fn join_comma_separated<'fmt>(
+        &'fmt mut self,
+        sequence_end: TextSize,
+    ) -> JoinCommaSeparatedBuilder<'fmt, 'ast, 'buf> {
+        JoinCommaSeparatedBuilder::new(self, sequence_end)
     }
 }
 
@@ -194,18 +203,20 @@ pub(crate) struct JoinCommaSeparatedBuilder<'fmt, 'ast, 'buf> {
     result: FormatResult<()>,
     fmt: &'fmt mut PyFormatter<'ast, 'buf>,
     end_of_last_entry: Option<TextSize>,
+    sequence_end: TextSize,
     /// We need to track whether we have more than one entry since a sole entry doesn't get a
     /// magic trailing comma even when expanded
     len: usize,
 }
 
 impl<'fmt, 'ast, 'buf> JoinCommaSeparatedBuilder<'fmt, 'ast, 'buf> {
-    fn new(f: &'fmt mut PyFormatter<'ast, 'buf>) -> Self {
+    fn new(f: &'fmt mut PyFormatter<'ast, 'buf>, sequence_end: TextSize) -> Self {
         Self {
             fmt: f,
             result: Ok(()),
             end_of_last_entry: None,
             len: 0,
+            sequence_end,
         }
     }
 
@@ -236,7 +247,7 @@ impl<'fmt, 'ast, 'buf> JoinCommaSeparatedBuilder<'fmt, 'ast, 'buf> {
     where
         T: Ranged,
         F: Format<PyFormatContext<'ast>>,
-        I: Iterator<Item = (T, F)>,
+        I: IntoIterator<Item = (T, F)>,
     {
         for (node, content) in entries {
             self.entry(&node, &content);
@@ -248,7 +259,7 @@ impl<'fmt, 'ast, 'buf> JoinCommaSeparatedBuilder<'fmt, 'ast, 'buf> {
     pub(crate) fn nodes<'a, T, I>(&mut self, entries: I) -> &mut Self
     where
         T: Ranged + AsFormat<PyFormatContext<'ast>> + 'a,
-        I: Iterator<Item = &'a T>,
+        I: IntoIterator<Item = &'a T>,
     {
         for node in entries {
             self.entry(node, &node.format());
@@ -260,14 +271,26 @@ impl<'fmt, 'ast, 'buf> JoinCommaSeparatedBuilder<'fmt, 'ast, 'buf> {
     pub(crate) fn finish(&mut self) -> FormatResult<()> {
         self.result.and_then(|_| {
             if let Some(last_end) = self.end_of_last_entry.take() {
-                let magic_trailing_comma = self.fmt.options().magic_trailing_comma().is_respect()
-                    && matches!(
-                        first_non_trivia_token(last_end, self.fmt.context().source()),
-                        Some(Token {
-                            kind: TokenKind::Comma,
-                            ..
-                        })
-                    );
+                let magic_trailing_comma = match self.fmt.options().magic_trailing_comma() {
+                    MagicTrailingComma::Respect => {
+                        let first_token = SimpleTokenizer::new(
+                            self.fmt.context().source(),
+                            TextRange::new(last_end, self.sequence_end),
+                        )
+                        .skip_trivia()
+                        // Skip over any closing parentheses belonging to the expression
+                        .find(|token| token.kind() != TokenKind::RParen);
+
+                        matches!(
+                            first_token,
+                            Some(Token {
+                                kind: TokenKind::Comma,
+                                ..
+                            })
+                        )
+                    }
+                    MagicTrailingComma::Ignore => false,
+                };
 
                 // If there is a single entry, only keep the magic trailing comma, don't add it if
                 // it wasn't there. If there is more than one entry, always add it.
@@ -287,13 +310,15 @@ impl<'fmt, 'ast, 'buf> JoinCommaSeparatedBuilder<'fmt, 'ast, 'buf> {
 
 #[cfg(test)]
 mod tests {
+    use rustpython_parser::ast::ModModule;
+    use rustpython_parser::Parse;
+
+    use ruff_formatter::format;
+
     use crate::comments::Comments;
     use crate::context::{NodeLevel, PyFormatContext};
     use crate::prelude::*;
     use crate::PyFormatOptions;
-    use ruff_formatter::format;
-    use rustpython_parser::ast::ModModule;
-    use rustpython_parser::Parse;
 
     fn format_ranged(level: NodeLevel) -> String {
         let source = r#"

crates/ruff_python_formatter/src/expression/expr_attribute.rs

@@ -1,11 +1,12 @@
-use crate::comments::{leading_comments, trailing_comments, Comments};
-use crate::expression::parentheses::{
-    default_expression_needs_parentheses, NeedsParentheses, Parentheses, Parenthesize,
-};
+use rustpython_parser::ast::{Constant, Expr, ExprAttribute, ExprConstant};
+
+use ruff_formatter::write;
+use ruff_python_ast::node::AnyNodeRef;
+
+use crate::comments::{leading_comments, trailing_comments};
+use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses, Parentheses};
 use crate::prelude::*;
 use crate::FormatNodeRule;
-use ruff_formatter::write;
-use rustpython_parser::ast::{Constant, Expr, ExprAttribute, ExprConstant};
 
 #[derive(Default)]
 pub struct FormatExprAttribute;
@@ -35,7 +36,7 @@ impl FormatNodeRule<ExprAttribute> for FormatExprAttribute {
             dangling_comments.split_at(leading_attribute_comments_start);
 
         if needs_parentheses {
-            value.format().with_options(Parenthesize::Always).fmt(f)?;
+            value.format().with_options(Parentheses::Always).fmt(f)?;
         } else if let Expr::Attribute(expr_attribute) = value.as_ref() {
             // We're in a attribute chain (`a.b.c`). The outermost node adds parentheses if
             // required, the inner ones don't need them so we skip the `Expr` formatting that
@@ -71,41 +72,23 @@ impl FormatNodeRule<ExprAttribute> for FormatExprAttribute {
     }
 }
 
-/// Checks if there are any own line comments in an attribute chain (a.b.c). This method is
-/// recursive up to the innermost expression that the attribute chain starts behind.
-fn has_breaking_comments_attribute_chain(
-    expr_attribute: &ExprAttribute,
-    comments: &Comments,
-) -> bool {
-    if comments
-        .dangling_comments(expr_attribute)
-        .iter()
-        .any(|comment| comment.line_position().is_own_line())
-        || comments.has_trailing_own_line_comments(expr_attribute)
-    {
-        return true;
-    }
-
-    if let Expr::Attribute(inner) = expr_attribute.value.as_ref() {
-        return has_breaking_comments_attribute_chain(inner, comments);
-    }
-
-    return comments.has_trailing_own_line_comments(expr_attribute.value.as_ref());
-}
-
 impl NeedsParentheses for ExprAttribute {
     fn needs_parentheses(
         &self,
-        parenthesize: Parenthesize,
+        parent: AnyNodeRef,
         context: &PyFormatContext,
-    ) -> Parentheses {
-        if has_breaking_comments_attribute_chain(self, context.comments()) {
-            return Parentheses::Always;
-        }
-
-        match default_expression_needs_parentheses(self.into(), parenthesize, context) {
-            Parentheses::Optional => Parentheses::Never,
-            parentheses => parentheses,
+    ) -> OptionalParentheses {
+        // Checks if there are any own line comments in an attribute chain (a.b.c).
+        if context
+            .comments()
+            .dangling_comments(self)
+            .iter()
+            .any(|comment| comment.line_position().is_own_line())
+            || context.comments().has_trailing_own_line_comments(self)
+        {
+            OptionalParentheses::Always
+        } else {
+            self.value.needs_parentheses(parent, context)
         }
     }
 }

crates/ruff_python_formatter/src/expression/expr_await.rs

@@ -1,10 +1,9 @@
 use crate::context::PyFormatContext;
-use crate::expression::parentheses::{
-    default_expression_needs_parentheses, NeedsParentheses, Parentheses, Parenthesize,
-};
+use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses};
 use crate::{AsFormat, FormatNodeRule, PyFormatter};
 use ruff_formatter::prelude::{space, text};
 use ruff_formatter::{write, Buffer, FormatResult};
+use ruff_python_ast::node::AnyNodeRef;
 use rustpython_parser::ast::ExprAwait;
 
 #[derive(Default)]
@@ -20,9 +19,9 @@ impl FormatNodeRule<ExprAwait> for FormatExprAwait {
 impl NeedsParentheses for ExprAwait {
     fn needs_parentheses(
         &self,
-        parenthesize: Parenthesize,
-        context: &PyFormatContext,
-    ) -> Parentheses {
-        default_expression_needs_parentheses(self.into(), parenthesize, context)
+        _parent: AnyNodeRef,
+        _context: &PyFormatContext,
+    ) -> OptionalParentheses {
+        OptionalParentheses::Multiline
     }
 }

crates/ruff_python_formatter/src/expression/expr_bin_op.rs

@@ -1,13 +1,12 @@
 use crate::comments::{trailing_comments, trailing_node_comments};
 use crate::expression::parentheses::{
-    default_expression_needs_parentheses, in_parentheses_only_group, is_expression_parenthesized,
-    NeedsParentheses, Parenthesize,
+    in_parentheses_only_group, is_expression_parenthesized, NeedsParentheses, OptionalParentheses,
 };
 use crate::expression::Parentheses;
 use crate::prelude::*;
 use crate::FormatNodeRule;
 use ruff_formatter::{write, FormatOwnedWithRule, FormatRefWithRule, FormatRuleWithOptions};
-use ruff_python_ast::node::AstNode;
+use ruff_python_ast::node::{AnyNodeRef, AstNode};
 use rustpython_parser::ast::{
     Constant, Expr, ExprAttribute, ExprBinOp, ExprConstant, ExprUnaryOp, Operator, UnaryOp,
 };
@@ -175,9 +174,9 @@ impl FormatRule<Operator, PyFormatContext<'_>> for FormatOperator {
 impl NeedsParentheses for ExprBinOp {
     fn needs_parentheses(
         &self,
-        parenthesize: Parenthesize,
-        context: &PyFormatContext,
-    ) -> Parentheses {
-        default_expression_needs_parentheses(self.into(), parenthesize, context)
+        _parent: AnyNodeRef,
+        _context: &PyFormatContext,
+    ) -> OptionalParentheses {
+        OptionalParentheses::Multiline
     }
 }

crates/ruff_python_formatter/src/expression/expr_bool_op.rs

@@ -1,10 +1,10 @@
 use crate::comments::leading_comments;
 use crate::expression::parentheses::{
-    default_expression_needs_parentheses, in_parentheses_only_group, NeedsParentheses, Parentheses,
-    Parenthesize,
+    in_parentheses_only_group, NeedsParentheses, OptionalParentheses, Parentheses,
 };
 use crate::prelude::*;
 use ruff_formatter::{write, FormatOwnedWithRule, FormatRefWithRule, FormatRuleWithOptions};
+use ruff_python_ast::node::AnyNodeRef;
 use rustpython_parser::ast::{BoolOp, ExprBoolOp};
 
 #[derive(Default)]
@@ -70,10 +70,10 @@ impl FormatNodeRule<ExprBoolOp> for FormatExprBoolOp {
 impl NeedsParentheses for ExprBoolOp {
     fn needs_parentheses(
         &self,
-        parenthesize: Parenthesize,
-        context: &PyFormatContext,
-    ) -> Parentheses {
-        default_expression_needs_parentheses(self.into(), parenthesize, context)
+        _parent: AnyNodeRef,
+        _context: &PyFormatContext,
+    ) -> OptionalParentheses {
+        OptionalParentheses::Multiline
     }
 }