[flake8-comprehensions] Handled special case for C401 which also matches C416

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

@@ -1,20 +1,30 @@
-x = set(x for x in range(3))
-x = set(x for x in range(3))
-y = f"{set(a if a < 6 else 0  for a in range(3))}"
-_ = "{}".format(set(a if a < 6 else 0 for a in range(3)))
-print(f"Hello {set(a for a in range(3))} World")
-
+# Cannot conbime with C416. Should use set comprehension here.
+even_nums = set(2 * x for x in range(3))
+odd_nums = set(
+    2 * x + 1 for x in range(3)
+)
+small_nums = f"{set(a if a < 6 else 0 for a in range(3))}"
 
 def f(x):
     return x
 
-
-print(f'Hello {set(a for a in "abc")} World')
-print(f"Hello {set(a for a in 'abc')} World")
 print(f"Hello {set(f(a) for a in 'abc')} World")
+print(f"Hello { set(f(a) for a in 'abc') } World")
+
+
+# Short-circuit case, combine with C416 and should produce x = set(range(3))
+x = set(x for x in range(3))
+x = set(
+    x for x in range(3)
+)
+print(f"Hello {set(a for a in range(3))} World")
 print(f"{set(a for a in 'abc') - set(a for a in 'ab')}")
 print(f"{ set(a for a in 'abc') - set(a for a in 'ab') }")
 
-# The fix generated for this diagnostic is incorrect, as we add additional space
-# around the set comprehension.
-print(f"{ {set(a for a in 'abc')} }")
+
+# Not built-in set.
+def set(*args, **kwargs):
+    return None
+
+set(2 * x for x in range(3))
+set(x for x in range(3))

crates/ruff_linter/src/rules/flake8_comprehensions/rules/unnecessary_generator_set.rs

@@ -1,6 +1,8 @@
 use ruff_diagnostics::{AlwaysFixableViolation, Diagnostic, Edit, Fix};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast as ast;
+use ruff_python_ast::comparable::ComparableExpr;
+use ruff_python_ast::ExprGenerator;
 use ruff_text_size::{Ranged, TextSize};
 
 use crate::checkers::ast::Checker;
@@ -10,37 +12,53 @@ use super::helpers;
 
 /// ## What it does
 /// Checks for unnecessary generators that can be rewritten as `set`
-/// comprehensions.
+/// comprehensions (or with `set` directly).
 ///
 /// ## Why is this bad?
 /// It is unnecessary to use `set` around a generator expression, since
 /// there are equivalent comprehensions for these types. Using a
 /// comprehension is clearer and more idiomatic.
 ///
+/// Further, if the comprehension can be removed entirely, as in the case of
+/// `set(x for x in foo)`, it's better to use `set(foo)` directly, since it's
+/// even more direct.
+///
 /// ## Examples
 /// ```python
 /// set(f(x) for x in foo)
+/// set(x for x in foo)
 /// ```
 ///
 /// Use instead:
 /// ```python
 /// {f(x) for x in foo}
+/// set(foo)
 /// ```
 ///
 /// ## Fix safety
 /// This rule's fix is marked as unsafe, as it may occasionally drop comments
 /// when rewriting the call. In most cases, though, comments will be preserved.
 #[violation]
-pub struct UnnecessaryGeneratorSet;
+pub struct UnnecessaryGeneratorSet {
+    short_circuit: bool,
+}
 
 impl AlwaysFixableViolation for UnnecessaryGeneratorSet {
     #[derive_message_formats]
     fn message(&self) -> String {
-        format!("Unnecessary generator (rewrite as a `set` comprehension)")
+        if self.short_circuit {
+            format!("Unnecessary generator (rewrite using `set()`")
+        } else {
+            format!("Unnecessary generator (rewrite as a `set` comprehension)")
+        }
     }
 
     fn fix_title(&self) -> String {
-        "Rewrite as a `set` comprehension".to_string()
+        if self.short_circuit {
+            "Rewrite using `set()`".to_string()
+        } else {
+            "Rewrite as a `set` comprehension".to_string()
+        }
     }
 }
 
@@ -57,28 +75,59 @@ pub(crate) fn unnecessary_generator_set(checker: &mut Checker, call: &ast::ExprC
     if !checker.semantic().is_builtin("set") {
         return;
     }
-    if argument.is_generator_expr() {
-        let mut diagnostic = Diagnostic::new(UnnecessaryGeneratorSet, call.range());
 
-        // Convert `set(x for x in y)` to `{x for x in y}`.
-        diagnostic.set_fix({
-            // Replace `set(` with `}`.
-            let call_start = Edit::replacement(
-                pad_start("{", call.range(), checker.locator(), checker.semantic()),
-                call.start(),
-                call.arguments.start() + TextSize::from(1),
-            );
+    let Some(ExprGenerator {
+        elt, generators, ..
+    }) = argument.as_generator_expr()
+    else {
+        return;
+    };
+
+    // Short-circuit: given `set(x for x in y)`, generate `set(y)` (in lieu of `{x for x in y}`).
+    if let [generator] = generators.as_slice() {
+        if generator.ifs.is_empty() && !generator.is_async {
+            if ComparableExpr::from(elt) == ComparableExpr::from(&generator.target) {
+                let mut diagnostic = Diagnostic::new(
+                    UnnecessaryGeneratorSet {
+                        short_circuit: true,
+                    },
+                    call.range(),
+                );
+                let iterator = format!("set({})", checker.locator().slice(generator.iter.range()));
+                diagnostic.set_fix(Fix::unsafe_edit(Edit::range_replacement(
+                    iterator,
+                    call.range(),
+                )));
+                checker.diagnostics.push(diagnostic);
+                return;
+            }
+        }
+    }
+
+    // Convert `set(f(x) for x in y)` to `{f(x) for x in y}`.
+    let mut diagnostic = Diagnostic::new(
+        UnnecessaryGeneratorSet {
+            short_circuit: false,
+        },
+        call.range(),
+    );
+    diagnostic.set_fix({
+        // Replace `set(` with `}`.
+        let call_start = Edit::replacement(
+            pad_start("{", call.range(), checker.locator(), checker.semantic()),
+            call.start(),
+            call.arguments.start() + TextSize::from(1),
+        );
 
-            // Replace `)` with `}`.
-            let call_end = Edit::replacement(
-                pad_end("}", call.range(), checker.locator(), checker.semantic()),
-                call.arguments.end() - TextSize::from(1),
-                call.end(),
-            );
+        // Replace `)` with `}`.
+        let call_end = Edit::replacement(
+            pad_end("}", call.range(), checker.locator(), checker.semantic()),
+            call.arguments.end() - TextSize::from(1),
+            call.end(),
+        );
 
-            Fix::unsafe_edits(call_start, [call_end])
-        });
+        Fix::unsafe_edits(call_start, [call_end])
+    });
 
-        checker.diagnostics.push(diagnostic);
-    }
+    checker.diagnostics.push(diagnostic);
 }