Implement C401 (#343)

README.md

@@ -276,6 +276,7 @@ The 🛠 emoji indicates that a rule is automatically fixable by the `--fix` com
 | A002 | BuiltinArgumentShadowing | Argument `...` is shadowing a python builtin |  |  |
 | A003 | BuiltinAttributeShadowing | Class attribute `...` is shadowing a python builtin |  |  |
 | C400 | UnnecessaryGeneratorList | Unnecessary generator - rewrite as a list comprehension |  |  |
+| C401 | UnnecessaryGeneratorSet | Unnecessary generator - rewrite as a set comprehension |  |  |
 | C403 | UnnecessaryListComprehensionSet | Unnecessary list comprehension - rewrite as a set comprehension |  |  |
 | C404 | UnnecessaryListComprehensionDict | Unnecessary list comprehension - rewrite as a dict comprehension |  |  |
 | SPR001 | SuperCallWithParameters | Use `super()` instead of `super(__class__, self)` |  | 🛠 |

resources/test/fixtures/C401.py

@@ -0,0 +1 @@
+x = set(x for x in range(3))

src/ast/checks.rs

@@ -731,6 +731,23 @@ pub fn unnecessary_generator_list(expr: &Expr, func: &Expr, args: &Vec<Expr>) ->
     None
 }
 
+/// Check `set(generator)` compliance.
+pub fn unnecessary_generator_set(expr: &Expr, func: &Expr, args: &Vec<Expr>) -> Option<Check> {
+    if args.len() == 1 {
+        if let ExprKind::Name { id, .. } = &func.node {
+            if id == "set" {
+                if let ExprKind::GeneratorExp { .. } = &args[0].node {
+                    return Some(Check::new(
+                        CheckKind::UnnecessaryGeneratorList,
+                        Range::from_located(expr),
+                    ));
+                }
+            }
+        }
+    }
+    None
+}
+
 /// Check `set([...])` compliance.
 pub fn unnecessary_list_comprehension_set(
     expr: &Expr,

src/check_ast.rs

@@ -772,6 +772,12 @@ where
                     };
                 }
 
+                if self.settings.enabled.contains(&CheckCode::C401) {
+                    if let Some(check) = checks::unnecessary_generator_set(expr, func, args) {
+                        self.checks.push(check);
+                    };
+                }
+
                 if self.settings.enabled.contains(&CheckCode::C403) {
                     if let Some(check) =
                         checks::unnecessary_list_comprehension_set(expr, func, args)

src/checks.rs

@@ -53,7 +53,7 @@ pub const DEFAULT_CHECK_CODES: [CheckCode; 42] = [
     CheckCode::F901,
 ];
 
-pub const ALL_CHECK_CODES: [CheckCode; 56] = [
+pub const ALL_CHECK_CODES: [CheckCode; 57] = [
     // pycodestyle
     CheckCode::E402,
     CheckCode::E501,
@@ -104,6 +104,7 @@ pub const ALL_CHECK_CODES: [CheckCode; 56] = [
     CheckCode::A003,
     // flake8-comprehensions
     CheckCode::C400,
+    CheckCode::C401,
     CheckCode::C403,
     CheckCode::C404,
     // flake8-super
@@ -173,6 +174,7 @@ pub enum CheckCode {
     A003,
     // flake8-comprehensions
     C400,
+    C401,
     C403,
     C404,
     // flake8-super
@@ -317,6 +319,7 @@ impl CheckCode {
             CheckCode::A003 => "A003",
             // flake8-comprehensions
             CheckCode::C400 => "C400",
+            CheckCode::C401 => "C401",
             CheckCode::C403 => "C403",
             CheckCode::C404 => "C404",
             // flake8-super
@@ -399,6 +402,7 @@ impl CheckCode {
             CheckCode::A003 => CheckKind::BuiltinAttributeShadowing("...".to_string()),
             // flake8-comprehensions
             CheckCode::C400 => CheckKind::UnnecessaryGeneratorList,
+            CheckCode::C401 => CheckKind::UnnecessaryGeneratorSet,
             CheckCode::C403 => CheckKind::UnnecessaryListComprehensionSet,
             CheckCode::C404 => CheckKind::UnnecessaryListComprehensionDict,
             // flake8-super
@@ -481,6 +485,7 @@ pub enum CheckKind {
     BuiltinAttributeShadowing(String),
     // flakes8-comprehensions
     UnnecessaryGeneratorList,
+    UnnecessaryGeneratorSet,
     UnnecessaryListComprehensionSet,
     UnnecessaryListComprehensionDict,
     // flake8-super
@@ -550,6 +555,7 @@ impl CheckKind {
             CheckKind::BuiltinAttributeShadowing(_) => "BuiltinAttributeShadowing",
             // flake8-comprehensions
             CheckKind::UnnecessaryGeneratorList => "UnnecessaryGeneratorList",
+            CheckKind::UnnecessaryGeneratorSet => "UnnecessaryGeneratorSet",
             CheckKind::UnnecessaryListComprehensionSet => "UnnecessaryListComprehensionSet",
             CheckKind::UnnecessaryListComprehensionDict => "UnnecessaryListComprehensionDict",
             // flake8-super
@@ -619,6 +625,7 @@ impl CheckKind {
             CheckKind::BuiltinAttributeShadowing(_) => &CheckCode::A003,
             // flake8-comprehensions
             CheckKind::UnnecessaryGeneratorList => &CheckCode::C400,
+            CheckKind::UnnecessaryGeneratorSet => &CheckCode::C401,
             CheckKind::UnnecessaryListComprehensionSet => &CheckCode::C403,
             CheckKind::UnnecessaryListComprehensionDict => &CheckCode::C404,
             // flake8-super
@@ -779,6 +786,9 @@ impl CheckKind {
             CheckKind::UnnecessaryGeneratorList => {
                 "Unnecessary generator - rewrite as a list comprehension".to_string()
             }
+            CheckKind::UnnecessaryGeneratorSet => {
+                "Unnecessary generator - rewrite as a set comprehension".to_string()
+            }
             CheckKind::UnnecessaryListComprehensionSet => {
                 "Unnecessary list comprehension - rewrite as a set comprehension".to_string()
             }

src/linter.rs

@@ -798,6 +798,18 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn c401() -> Result<()> {
+        let mut checks = check_path(
+            Path::new("./resources/test/fixtures/C401.py"),
+            &settings::Settings::for_rule(CheckCode::C401),
+            &fixer::Mode::Generate,
+        )?;
+        checks.sort_by_key(|check| check.location);
+        insta::assert_yaml_snapshot!(checks);
+        Ok(())
+    }
+
     #[test]
     fn c403() -> Result<()> {
         let mut checks = check_path(

src/snapshots/ruff__linter__tests__c401.snap

@@ -0,0 +1,13 @@
+---
+source: src/linter.rs
+expression: checks
+---
+- kind: UnnecessaryGeneratorList
+  location:
+    row: 1
+    column: 5
+  end_location:
+    row: 1
+    column: 29
+  fix: ~
+