[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 <dylwil3@gmail.com>
This commit is contained in:
Ayush Baweja 2025-02-15 12:45:41 -05:00 committed by GitHub
parent 3c69b685ee
commit df45a9db64
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 51 additions and 14 deletions

View file

@ -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)],)

View file

@ -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());

View file

@ -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()),

View file

@ -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(),
);

View file

@ -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)}