Extend unncessary-generator-any-all to sets

astral-sh · Mar 30, 2023 · 6c12076 · 6c12076
1 parent 8fb9b8c
commit 6c12076
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 25 deletions.
diff --git a/crates/ruff/resources/test/fixtures/flake8_pie/PIE802.py b/crates/ruff/resources/test/fixtures/flake8_pie/PIE802.py
@@ -7,14 +7,13 @@
 all(  # first comment
     [x.id for x in bar],  # second comment
 )  # third comment
+any({x.id for x in bar})
 
 # OK
-all((x.id for x in bar))
 all(x.id for x in bar)
 all(x.id for x in bar)
 any(x.id for x in bar)
-any({x.id for x in bar})
-
+all((x.id for x in bar))
 
 async def f() -> bool:
     return all([await use_greeting(greeting) for greeting in await greetings()])
diff --git a/crates/ruff/src/rules/flake8_pie/rules.rs b/crates/ruff/src/rules/flake8_pie/rules.rs
@@ -344,29 +344,26 @@ pub fn unnecessary_comprehension_any_all(
     func: &Expr,
     args: &[Expr],
 ) {
-    if let ExprKind::Name { id, .. } = &func.node {
-        if (id == "all" || id == "any") && args.len() == 1 {
-            if !checker.ctx.is_builtin(id) {
-                return;
-            }
-            if let ExprKind::ListComp { elt, .. } = &args[0].node {
-                if is_async_generator(elt) {
-                    return;
-                }
-                let mut diagnostic =
-                    Diagnostic::new(UnnecessaryComprehensionAnyAll, Range::from(&args[0]));
-                if checker.patch(diagnostic.kind.rule()) {
-                    diagnostic.try_set_fix(|| {
-                        fixes::fix_unnecessary_comprehension_any_all(
-                            checker.locator,
-                            checker.stylist,
-                            expr,
-                        )
-                    });
-                }
-                checker.diagnostics.push(diagnostic);
-            }
+    let ExprKind::Name { id, .. } = &func.node  else {
+        return;
+    };
+    if (id == "all" || id == "any") && args.len() == 1 {
+        let (ExprKind::ListComp { elt, .. } | ExprKind::SetComp { elt, .. }) = &args[0].node else {
+            return;
+        };
+        if is_async_generator(elt) {
+            return;
+        }
+        if !checker.ctx.is_builtin(id) {
+            return;
+        }
+        let mut diagnostic = Diagnostic::new(UnnecessaryComprehensionAnyAll, Range::from(&args[0]));
+        if checker.patch(diagnostic.kind.rule()) {
+            diagnostic.try_set_fix(|| {
+                fixes::fix_unnecessary_comprehension_any_all(checker.locator, checker.stylist, expr)
+            });
         }
+        checker.diagnostics.push(diagnostic);
     }
 }
 

diff --git a/...ruff/src/rules/flake8_pie/snapshots/ruff__rules__flake8_pie__tests__PIE802_PIE802.py.snap b/...ruff/src/rules/flake8_pie/snapshots/ruff__rules__flake8_pie__tests__PIE802_PIE802.py.snap
@@ -86,4 +86,18 @@ expression: diagnostics
           row: 9
           column: 1
   parent: ~
+- kind:
+    name: UnnecessaryComprehensionAnyAll
+    body: Unnecessary list comprehension.
+    suggestion: Remove unnecessary list comprehension
+    fixable: true
+  location:
+    row: 10
+    column: 4
+  end_location:
+    row: 10
+    column: 23
+  fix:
+    edits: []
+  parent: ~