Add parenthesized flag to ExprTuple and ExprGenerator (#9614)

crates/ruff_python_formatter/src/comments/placement.rs

@@ -12,7 +12,6 @@ use ruff_source_file::Locator;
 use ruff_text_size::{Ranged, TextLen, TextRange, TextSize};
 
 use crate::comments::visitor::{CommentPlacement, DecoratedComment};
-use crate::expression::expr_generator_exp::is_generator_parenthesized;
 use crate::expression::expr_slice::{assign_comment_in_slice, ExprSliceCommentSection};
 use crate::other::parameters::{
     assign_argument_separator_comment_placement, find_parameter_separators,
@@ -315,12 +314,11 @@ fn handle_enclosed_comment<'a>(
         | AnyNodeRef::ExprSet(_)
         | AnyNodeRef::ExprListComp(_)
         | AnyNodeRef::ExprSetComp(_) => handle_bracketed_end_of_line_comment(comment, locator),
-        AnyNodeRef::ExprTuple(tuple) if tuple.is_parenthesized(locator.contents()) => {
-            handle_bracketed_end_of_line_comment(comment, locator)
-        }
-        AnyNodeRef::ExprGeneratorExp(generator)
-            if is_generator_parenthesized(generator, locator.contents()) =>
-        {
+        AnyNodeRef::ExprTuple(ast::ExprTuple {
+            parenthesized: true,
+            ..
+        }) => handle_bracketed_end_of_line_comment(comment, locator),
+        AnyNodeRef::ExprGeneratorExp(generator) if generator.parenthesized => {
             handle_bracketed_end_of_line_comment(comment, locator)
         }
         _ => CommentPlacement::Default(comment),

crates/ruff_python_formatter/src/expression/expr_generator_exp.rs

@@ -1,8 +1,6 @@
 use ruff_formatter::{format_args, write, FormatRuleWithOptions};
 use ruff_python_ast::AnyNodeRef;
 use ruff_python_ast::ExprGeneratorExp;
-use ruff_python_trivia::{SimpleTokenKind, SimpleTokenizer};
-use ruff_text_size::{Ranged, TextRange};
 
 use crate::comments::SourceComment;
 use crate::expression::parentheses::{parenthesized, NeedsParentheses, OptionalParentheses};
@@ -42,6 +40,7 @@ impl FormatNodeRule<ExprGeneratorExp> for FormatExprGeneratorExp {
             range: _,
             elt,
             generators,
+            parenthesized: is_parenthesized,
         } = item;
 
         let joined = format_with(|f| {
@@ -55,7 +54,7 @@ impl FormatNodeRule<ExprGeneratorExp> for FormatExprGeneratorExp {
 
         if self.parentheses == GeneratorExpParentheses::Preserve
             && dangling.is_empty()
-            && !is_generator_parenthesized(item, f.context().source())
+            && !is_parenthesized
         {
             write!(
                 f,
@@ -101,37 +100,3 @@ impl NeedsParentheses for ExprGeneratorExp {
         }
     }
 }
-
-/// Return `true` if a generator is parenthesized in the source code.
-pub(crate) fn is_generator_parenthesized(generator: &ExprGeneratorExp, source: &str) -> bool {
-    // Count the number of open parentheses between the start of the generator and the first element.
-    let open_parentheses_count = SimpleTokenizer::new(
-        source,
-        TextRange::new(generator.start(), generator.elt.start()),
-    )
-    .skip_trivia()
-    .filter(|token| token.kind() == SimpleTokenKind::LParen)
-    .count();
-    if open_parentheses_count == 0 {
-        return false;
-    }
-
-    // Count the number of parentheses between the end of the generator and its trailing comma.
-    let close_parentheses_count = SimpleTokenizer::new(
-        source,
-        TextRange::new(
-            generator.elt.end(),
-            generator
-                .generators
-                .first()
-                .map_or(generator.end(), Ranged::start),
-        ),
-    )
-    .skip_trivia()
-    .filter(|token| token.kind() == SimpleTokenKind::RParen)
-    .count();
-
-    // If the number of open parentheses is greater than the number of close parentheses, the
-    // generator is parenthesized.
-    open_parentheses_count > close_parentheses_count
-}

crates/ruff_python_formatter/src/expression/expr_tuple.rs

@@ -116,6 +116,7 @@ impl FormatNodeRule<ExprTuple> for FormatExprTuple {
             elts,
             ctx: _,
             range: _,
+            parenthesized: is_parenthesized,
         } = item;
 
         let comments = f.context().comments().clone();
@@ -136,7 +137,7 @@ impl FormatNodeRule<ExprTuple> for FormatExprTuple {
                 return empty_parenthesized("(", dangling, ")").fmt(f);
             }
             [single] => match self.parentheses {
-                TupleParentheses::Preserve if !item.is_parenthesized(f.context().source()) => {
+                TupleParentheses::Preserve if !is_parenthesized => {
                     write!(f, [single.format(), token(",")])
                 }
                 _ =>
@@ -152,7 +153,7 @@ impl FormatNodeRule<ExprTuple> for FormatExprTuple {
             //
             // Unlike other expression parentheses, tuple parentheses are part of the range of the
             // tuple itself.
-            _ if item.is_parenthesized(f.context().source())
+            _ if *is_parenthesized
                 && !(self.parentheses == TupleParentheses::NeverPreserve
                     && dangling.is_empty()) =>
             {

crates/ruff_python_formatter/src/expression/mod.rs

@@ -14,7 +14,6 @@ use ruff_text_size::Ranged;
 use crate::builders::parenthesize_if_expands;
 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::parentheses::{
     is_expression_parenthesized, optional_parentheses, parenthesized, NeedsParentheses,
     OptionalParentheses, Parentheses, Parenthesize,
@@ -661,15 +660,16 @@ impl<'input> CanOmitOptionalParenthesesVisitor<'input> {
                 return;
             }
 
-            Expr::Tuple(tuple) if tuple.is_parenthesized(self.context.source()) => {
+            Expr::Tuple(ast::ExprTuple {
+                parenthesized: true,
+                ..
+            }) => {
                 self.any_parenthesized_expressions = true;
                 // The values are always parenthesized, don't visit.
                 return;
             }
 
-            Expr::GeneratorExp(generator)
-                if is_generator_parenthesized(generator, self.context.source()) =>
-            {
+            Expr::GeneratorExp(generator) if generator.parenthesized => {
                 self.any_parenthesized_expressions = true;
                 // The values are always parenthesized, don't visit.
                 return;
@@ -1035,11 +1035,7 @@ pub(crate) fn has_own_parentheses(
             Some(OwnParentheses::NonEmpty)
         }
 
-        Expr::GeneratorExp(generator)
-            if is_generator_parenthesized(generator, context.source()) =>
-        {
-            Some(OwnParentheses::NonEmpty)
-        }
+        Expr::GeneratorExp(generator) if generator.parenthesized => Some(OwnParentheses::NonEmpty),
 
         // These expressions must contain _some_ child or trivia token in order to be non-empty.
         Expr::List(ast::ExprList { elts, .. }) | Expr::Set(ast::ExprSet { elts, .. }) => {
@@ -1050,7 +1046,12 @@ pub(crate) fn has_own_parentheses(
             }
         }
 
-        Expr::Tuple(tuple) if tuple.is_parenthesized(context.source()) => {
+        Expr::Tuple(
+            tuple @ ast::ExprTuple {
+                parenthesized: true,
+                ..
+            },
+        ) => {
             if !tuple.elts.is_empty() || context.comments().has_dangling(AnyNodeRef::from(expr)) {
                 Some(OwnParentheses::NonEmpty)
             } else {