From df45a9db641dc319184548cbf6685348657b10b8 Mon Sep 17 00:00:00 2001 From: Ayush Baweja <44344063+ayushbaweja@users.noreply.github.com> Date: Sat, 15 Feb 2025 12:45:41 -0500 Subject: [PATCH] [flake8-comprehensions]: Handle trailing comma in C403 fix (#16110) ## Summary Resolves [#16099 ](https://github.com/astral-sh/ruff/issues/16099) based on [#15929 ](https://github.com/astral-sh/ruff/pull/15929) ## Test Plan Added test case `s = set([x for x in range(3)],)` and updated snapshot. --------- Co-authored-by: dylwil3 --- .../fixtures/flake8_comprehensions/C403.py | 3 +++ .../rules/unnecessary_generator_list.rs | 14 +++++++------ .../rules/unnecessary_generator_set.rs | 14 +++++++------ .../unnecessary_list_comprehension_set.rs | 14 +++++++++++-- ...8_comprehensions__tests__C403_C403.py.snap | 20 +++++++++++++++++++ 5 files changed, 51 insertions(+), 14 deletions(-) diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_comprehensions/C403.py b/crates/ruff_linter/resources/test/fixtures/flake8_comprehensions/C403.py index ba0d2598ea..c1a9feb24c 100644 --- a/crates/ruff_linter/resources/test/fixtures/flake8_comprehensions/C403.py +++ b/crates/ruff_linter/resources/test/fixtures/flake8_comprehensions/C403.py @@ -32,3 +32,6 @@ s = set( # outer set comment [ # comprehension comment x for x in range(3)] )))) + +# Test trailing comma case +s = set([x for x in range(3)],) \ No newline at end of file diff --git a/crates/ruff_linter/src/rules/flake8_comprehensions/rules/unnecessary_generator_list.rs b/crates/ruff_linter/src/rules/flake8_comprehensions/rules/unnecessary_generator_list.rs index cc389e1f4d..c2345fdcce 100644 --- a/crates/ruff_linter/src/rules/flake8_comprehensions/rules/unnecessary_generator_list.rs +++ b/crates/ruff_linter/src/rules/flake8_comprehensions/rules/unnecessary_generator_list.rs @@ -4,7 +4,7 @@ use ruff_python_ast as ast; use ruff_python_ast::comparable::ComparableExpr; use ruff_python_ast::parenthesize::parenthesized_range; use ruff_python_ast::ExprGenerator; -use ruff_python_trivia::{SimpleTokenKind, SimpleTokenizer}; +use ruff_python_parser::TokenKind; use ruff_text_size::{Ranged, TextRange, TextSize}; use crate::checkers::ast::Checker; @@ -125,11 +125,13 @@ pub(crate) fn unnecessary_generator_list(checker: &Checker, call: &ast::ExprCall // Replace `)` with `]`. // Place `]` at argument's end or at trailing comma if present - let mut tokenizer = - SimpleTokenizer::new(checker.source(), TextRange::new(argument.end(), call.end())); - let right_bracket_loc = tokenizer - .find(|token| token.kind == SimpleTokenKind::Comma) - .map_or(call.arguments.end(), |comma| comma.end()) + let after_arg_tokens = checker + .tokens() + .in_range(TextRange::new(argument.end(), call.end())); + let right_bracket_loc = after_arg_tokens + .iter() + .find(|token| token.kind() == TokenKind::Comma) + .map_or(call.arguments.end(), Ranged::end) - TextSize::from(1); let call_end = Edit::replacement("]".to_string(), right_bracket_loc, call.end()); diff --git a/crates/ruff_linter/src/rules/flake8_comprehensions/rules/unnecessary_generator_set.rs b/crates/ruff_linter/src/rules/flake8_comprehensions/rules/unnecessary_generator_set.rs index 193d975737..e04c67cb94 100644 --- a/crates/ruff_linter/src/rules/flake8_comprehensions/rules/unnecessary_generator_set.rs +++ b/crates/ruff_linter/src/rules/flake8_comprehensions/rules/unnecessary_generator_set.rs @@ -4,7 +4,7 @@ use ruff_python_ast as ast; use ruff_python_ast::comparable::ComparableExpr; use ruff_python_ast::parenthesize::parenthesized_range; use ruff_python_ast::ExprGenerator; -use ruff_python_trivia::{SimpleTokenKind, SimpleTokenizer}; +use ruff_python_parser::TokenKind; use ruff_text_size::{Ranged, TextRange, TextSize}; use crate::checkers::ast::Checker; @@ -128,11 +128,13 @@ pub(crate) fn unnecessary_generator_set(checker: &Checker, call: &ast::ExprCall) // Replace `)` with `}`. // Place `}` at argument's end or at trailing comma if present - let mut tokenizer = - SimpleTokenizer::new(checker.source(), TextRange::new(argument.end(), call.end())); - let right_brace_loc = tokenizer - .find(|token| token.kind == SimpleTokenKind::Comma) - .map_or(call.arguments.end(), |comma| comma.end()) + let after_arg_tokens = checker + .tokens() + .in_range(TextRange::new(argument.end(), call.end())); + let right_brace_loc = after_arg_tokens + .iter() + .find(|token| token.kind() == TokenKind::Comma) + .map_or(call.arguments.end(), Ranged::end) - TextSize::from(1); let call_end = Edit::replacement( pad_end("}", call.range(), checker.locator(), checker.semantic()), diff --git a/crates/ruff_linter/src/rules/flake8_comprehensions/rules/unnecessary_list_comprehension_set.rs b/crates/ruff_linter/src/rules/flake8_comprehensions/rules/unnecessary_list_comprehension_set.rs index 693e45f4c8..4a91ff7c5e 100644 --- a/crates/ruff_linter/src/rules/flake8_comprehensions/rules/unnecessary_list_comprehension_set.rs +++ b/crates/ruff_linter/src/rules/flake8_comprehensions/rules/unnecessary_list_comprehension_set.rs @@ -2,7 +2,8 @@ use ruff_diagnostics::{AlwaysFixableViolation, Diagnostic, Edit, Fix}; use ruff_macros::{derive_message_formats, ViolationMetadata}; use ruff_python_ast as ast; use ruff_python_ast::parenthesize::parenthesized_range; -use ruff_text_size::{Ranged, TextSize}; +use ruff_python_parser::TokenKind; +use ruff_text_size::{Ranged, TextRange, TextSize}; use crate::checkers::ast::Checker; use crate::rules::flake8_comprehensions::fixes::{pad_end, pad_start}; @@ -70,9 +71,18 @@ pub(crate) fn unnecessary_list_comprehension_set(checker: &Checker, call: &ast:: ); // Replace `)` with `}`. + // Place `}` at argument's end or at trailing comma if present + let after_arg_tokens = checker + .tokens() + .in_range(TextRange::new(argument.end(), call.end())); + let right_brace_loc = after_arg_tokens + .iter() + .find(|token| token.kind() == TokenKind::Comma) + .map_or(call.arguments.end() - one, |comma| comma.end() - one); + let call_end = Edit::replacement( pad_end("}", call.range(), checker.locator(), checker.semantic()), - call.arguments.end() - one, + right_brace_loc, call.end(), ); diff --git a/crates/ruff_linter/src/rules/flake8_comprehensions/snapshots/ruff_linter__rules__flake8_comprehensions__tests__C403_C403.py.snap b/crates/ruff_linter/src/rules/flake8_comprehensions/snapshots/ruff_linter__rules__flake8_comprehensions__tests__C403_C403.py.snap index fa588c4210..57b06b06d6 100644 --- a/crates/ruff_linter/src/rules/flake8_comprehensions/snapshots/ruff_linter__rules__flake8_comprehensions__tests__C403_C403.py.snap +++ b/crates/ruff_linter/src/rules/flake8_comprehensions/snapshots/ruff_linter__rules__flake8_comprehensions__tests__C403_C403.py.snap @@ -292,6 +292,8 @@ C403.py:29:5: C403 [*] Unnecessary list comprehension (rewrite as a set comprehe 33 | | x for x in range(3)] 34 | | )))) | |_____^ C403 +35 | +36 | # Test trailing comma case | = help: Rewrite as a set comprehension @@ -308,3 +310,21 @@ C403.py:29:5: C403 [*] Unnecessary list comprehension (rewrite as a set comprehe 29 |+s = { # outer set comment 30 |+ # comprehension comment 31 |+ x for x in range(3)} +35 32 | +36 33 | # Test trailing comma case +37 34 | s = set([x for x in range(3)],) + +C403.py:37:5: C403 [*] Unnecessary list comprehension (rewrite as a set comprehension) + | +36 | # Test trailing comma case +37 | s = set([x for x in range(3)],) + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^ C403 + | + = help: Rewrite as a set comprehension + +ℹ Unsafe fix +34 34 | )))) +35 35 | +36 36 | # Test trailing comma case +37 |-s = set([x for x in range(3)],) + 37 |+s = {x for x in range(3)}