From a352f2f092aee150cf78b6e6b940244f049b48d6 Mon Sep 17 00:00:00 2001 From: Micha Reiser Date: Fri, 8 Sep 2023 10:53:34 +0200 Subject: [PATCH] Preserve generator parentheses in single argument call expressions (#7226) --- .../fixtures/ruff/expression/generator_exp.py | 5 ++ .../src/expression/expr_generator_exp.rs | 48 ++++++++++++++++--- .../src/other/arguments.rs | 2 +- .../snapshots/format@expression__call.py.snap | 12 +++-- .../format@expression__generator_exp.py.snap | 14 +++++- 5 files changed, 67 insertions(+), 14 deletions(-) diff --git a/crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/generator_exp.py b/crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/generator_exp.py index 7baf9602e7..1fa75f6986 100644 --- a/crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/generator_exp.py +++ b/crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/generator_exp.py @@ -16,6 +16,11 @@ f((1) for _ in (a)) # combination of the two above f(((1) for _ in (a))) +bases = tuple( + (base._meta.label_lower if hasattr(base, "_meta") else base) + for base in flattened_bases +) + # black keeps these atm, but intends to remove them in the future: # https://github.com/psf/black/issues/2943 diff --git a/crates/ruff_python_formatter/src/expression/expr_generator_exp.rs b/crates/ruff_python_formatter/src/expression/expr_generator_exp.rs index 706102df39..75d2ba2e84 100644 --- a/crates/ruff_python_formatter/src/expression/expr_generator_exp.rs +++ b/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> for FormatExprGeneratorExp { @@ -51,8 +53,9 @@ impl FormatNodeRule 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 +} diff --git a/crates/ruff_python_formatter/src/other/arguments.rs b/crates/ruff_python_formatter/src/other/arguments.rs index f0c2651361..2f2335d3bc 100644 --- a/crates/ruff_python_formatter/src/other/arguments.rs +++ b/crates/ruff_python_formatter/src/other/arguments.rs @@ -37,7 +37,7 @@ impl FormatNodeRule for FormatArguments { generator_exp, &generator_exp .format() - .with_options(GeneratorExpParentheses::StripIfOnlyFunctionArg), + .with_options(GeneratorExpParentheses::Preserve), ), other => { let parentheses = diff --git a/crates/ruff_python_formatter/tests/snapshots/format@expression__call.py.snap b/crates/ruff_python_formatter/tests/snapshots/format@expression__call.py.snap index a8877d445c..9c26b8a3d5 100644 --- a/crates/ruff_python_formatter/tests/snapshots/format@expression__call.py.snap +++ b/crates/ruff_python_formatter/tests/snapshots/format@expression__call.py.snap @@ -386,10 +386,10 @@ threshold_date = datetime.datetime.now() - datetime.timedelta( # comment ) # Parenthesized and opening-parenthesis comments -func(x for x in y) +func((x for x in y)) func( # outer comment - x for x in y + (x for x in y) ) func( @@ -399,9 +399,11 @@ func( ) func( - # inner comment - x - for x in y + ( + # inner comment + x + for x in y + ) ) func( # outer comment diff --git a/crates/ruff_python_formatter/tests/snapshots/format@expression__generator_exp.py.snap b/crates/ruff_python_formatter/tests/snapshots/format@expression__generator_exp.py.snap index ed11b4724c..cc166a931b 100644 --- a/crates/ruff_python_formatter/tests/snapshots/format@expression__generator_exp.py.snap +++ b/crates/ruff_python_formatter/tests/snapshots/format@expression__generator_exp.py.snap @@ -22,6 +22,11 @@ f((1) for _ in (a)) # combination of the two above f(((1) for _ in (a))) +bases = tuple( + (base._meta.label_lower if hasattr(base, "_meta") else base) + for base in flattened_bases +) + # black keeps these atm, but intends to remove them in the future: # https://github.com/psf/black/issues/2943 @@ -67,13 +72,18 @@ sum((a for b in c), start=0) # black keeps these atm, but intends to remove them in the future: # https://github.com/psf/black/issues/2943 -f(1 for _ in a) +f((1 for _ in a)) # make sure source parenthesis detection isn't fooled by these f((1) for _ in (a)) # combination of the two above -f((1) for _ in (a)) +f(((1) for _ in (a))) + +bases = tuple( + (base._meta.label_lower if hasattr(base, "_meta") else base) + for base in flattened_bases +) # black keeps these atm, but intends to remove them in the future: