Preserve generator parentheses in single argument call expressions (#7226)

crates/ruff_python_formatter/src/expression/expr_generator_exp.rs

@@ -1,9 +1,10 @@
-use ruff_formatter::{format_args, write, Buffer, FormatResult, FormatRuleWithOptions};
+use ruff_formatter::{format_args, write, FormatRuleWithOptions};
 use ruff_python_ast::node::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};
 use crate::prelude::*;
 
@@ -12,12 +13,13 @@ pub enum GeneratorExpParentheses {
     #[default]
     Default,
 
-    /// Skip parens if the generator is the only argument to a function and doesn't contain any
-    /// dangling comments. For example:
+    /// Skips the parentheses if they aren't present in the source code. Used when formatting call expressions
+    /// because the parentheses are optional if the generator is the **only** argument:
+    ///
     /// ```python
     /// all(x for y in z)`
     /// ```
-    StripIfOnlyFunctionArg,
+    Preserve,
 }
 
 impl FormatRuleWithOptions<ExprGeneratorExp, PyFormatContext<'_>> for FormatExprGeneratorExp {
@@ -51,8 +53,9 @@ impl FormatNodeRule<ExprGeneratorExp> for FormatExprGeneratorExp {
         let comments = f.context().comments().clone();
         let dangling = comments.dangling(item);
 
-        if self.parentheses == GeneratorExpParentheses::StripIfOnlyFunctionArg
+        if self.parentheses == GeneratorExpParentheses::Preserve
             && dangling.is_empty()
+            && !is_generator_parenthesized(item, f.context().source())
         {
             write!(
                 f,
@@ -94,3 +97,36 @@ impl NeedsParentheses for ExprGeneratorExp {
         OptionalParentheses::Never
     }
 }
+
+fn is_generator_parenthesized(generator: &ExprGeneratorExp, source: &str) -> bool {
+    // / Count the number of open parentheses between the start of the tuple 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 first element 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/other/arguments.rs

@@ -37,7 +37,7 @@ impl FormatNodeRule<Arguments> for FormatArguments {
                             generator_exp,
                             &generator_exp
                                 .format()
-                                .with_options(GeneratorExpParentheses::StripIfOnlyFunctionArg),
+                                .with_options(GeneratorExpParentheses::Preserve),
                         ),
                         other => {
                             let parentheses =