[flake8-comprehensions]: Handle trailing comma in C403 fix (#16110)

## Summary

Resolves [#16099 ](#16099) based
on [#15929 ](#15929)

## Test Plan

Added test case `s = set([x for x in range(3)],)` and updated snapshot.

---------

Co-authored-by: dylwil3 <[email protected]>

Author: ayushbaweja

crates/ruff_linter/resources/test/fixtures/flake8_comprehensions/C403.py

@@ -32,3 +32,6 @@ def f(x):
 [ # comprehension comment
  x for x in range(3)]
  ))))
+
+# Test trailing comma case
+s = set([x for x in range(3)],)

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

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

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

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