Check SIM118 in comprehension (#1627)

astral-sh · Jan 4, 2023 · 0a0e192 · 0a0e192
1 parent 4473c7e
commit 0a0e192
Show file tree

Hide file tree

Showing 3 changed files with 96 additions and 1 deletion.
diff --git a/resources/test/fixtures/flake8_simplify/SIM118.py b/resources/test/fixtures/flake8_simplify/SIM118.py
@@ -12,3 +12,11 @@
 for key in list(obj.keys()):
     if some_property(key):
         del obj[key]
+
+[k for k in obj.keys()]  # SIM118
+
+{k for k in obj.keys()}  # SIM118
+
+{k: k for k in obj.keys()}  # SIM118
+
+(k for k in obj.keys())  # SIM118
diff --git a/src/checkers/ast.rs b/src/checkers/ast.rs
@@ -6,7 +6,7 @@ use itertools::Itertools;
 use log::error;
 use nohash_hasher::IntMap;
 use rustc_hash::{FxHashMap, FxHashSet};
-use rustpython_ast::{Located, Location};
+use rustpython_ast::{Comprehension, Located, Location};
 use rustpython_common::cformat::{CFormatError, CFormatErrorType};
 use rustpython_parser::ast::{
     Arg, Arguments, Constant, Excepthandler, ExcepthandlerKind, Expr, ExprContext, ExprKind,
@@ -2925,6 +2925,17 @@ where
         }
     }
 
+    fn visit_comprehension(&mut self, comprehension: &'b Comprehension) {
+        if self.settings.enabled.contains(&CheckCode::SIM118) {
+            flake8_simplify::plugins::key_in_dict_for(
+                self,
+                &comprehension.target,
+                &comprehension.iter,
+            );
+        }
+        visitor::walk_comprehension(self, comprehension);
+    }
+
     fn visit_arguments(&mut self, arguments: &'b Arguments) {
         if self.settings.enabled.contains(&CheckCode::B006) {
             flake8_bugbear::plugins::mutable_argument_default(self, arguments);

diff --git a/src/flake8_simplify/snapshots/ruff__flake8_simplify__tests__SIM118_SIM118.py.snap b/src/flake8_simplify/snapshots/ruff__flake8_simplify__tests__SIM118_SIM118.py.snap
@@ -97,4 +97,80 @@ expression: checks
       row: 9
       column: 21
   parent: ~
+- kind:
+    KeyInDict:
+      - k
+      - obj
+  location:
+    row: 16
+    column: 7
+  end_location:
+    row: 16
+    column: 22
+  fix:
+    content: obj
+    location:
+      row: 16
+      column: 12
+    end_location:
+      row: 16
+      column: 22
+  parent: ~
+- kind:
+    KeyInDict:
+      - k
+      - obj
+  location:
+    row: 18
+    column: 7
+  end_location:
+    row: 18
+    column: 22
+  fix:
+    content: obj
+    location:
+      row: 18
+      column: 12
+    end_location:
+      row: 18
+      column: 22
+  parent: ~
+- kind:
+    KeyInDict:
+      - k
+      - obj
+  location:
+    row: 20
+    column: 10
+  end_location:
+    row: 20
+    column: 25
+  fix:
+    content: obj
+    location:
+      row: 20
+      column: 15
+    end_location:
+      row: 20
+      column: 25
+  parent: ~
+- kind:
+    KeyInDict:
+      - k
+      - obj
+  location:
+    row: 22
+    column: 7
+  end_location:
+    row: 22
+    column: 22
+  fix:
+    content: obj
+    location:
+      row: 22
+      column: 12
+    end_location:
+      row: 22
+      column: 22
+  parent: ~