Move comprehension rules into shared analyze method (#6042)

crates/ruff/src/checkers/ast/analyze/comprehension.rs

@@ -0,0 +1,16 @@
+use rustpython_parser::ast::Comprehension;
+
+use crate::checkers::ast::Checker;
+use crate::codes::Rule;
+use crate::rules::flake8_simplify;
+
+/// Run lint rules over a [`Comprehension`] syntax nodes.
+pub(crate) fn comprehension(comprehension: &Comprehension, checker: &mut Checker) {
+    if checker.enabled(Rule::InDictKeys) {
+        flake8_simplify::rules::key_in_dict_for(
+            checker,
+            &comprehension.target,
+            &comprehension.iter,
+        );
+    }
+}

crates/ruff/src/checkers/ast/analyze/expression.rs

@@ -1305,15 +1305,6 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
             if checker.enabled(Rule::FunctionUsesLoopVariable) {
                 flake8_bugbear::rules::function_uses_loop_variable(checker, &Node::Expr(expr));
             }
-            if checker.enabled(Rule::InDictKeys) {
-                for generator in generators {
-                    flake8_simplify::rules::key_in_dict_for(
-                        checker,
-                        &generator.target,
-                        &generator.iter,
-                    );
-                }
-            }
             if checker.enabled(Rule::IterationOverSet) {
                 for generator in generators {
                     pylint::rules::iteration_over_set(checker, &generator.iter);
@@ -1334,15 +1325,6 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
             if checker.enabled(Rule::FunctionUsesLoopVariable) {
                 flake8_bugbear::rules::function_uses_loop_variable(checker, &Node::Expr(expr));
             }
-            if checker.enabled(Rule::InDictKeys) {
-                for generator in generators {
-                    flake8_simplify::rules::key_in_dict_for(
-                        checker,
-                        &generator.target,
-                        &generator.iter,
-                    );
-                }
-            }
             if checker.enabled(Rule::IterationOverSet) {
                 for generator in generators {
                     pylint::rules::iteration_over_set(checker, &generator.iter);
@@ -1360,15 +1342,6 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
             if checker.enabled(Rule::FunctionUsesLoopVariable) {
                 flake8_bugbear::rules::function_uses_loop_variable(checker, &Node::Expr(expr));
             }
-            if checker.enabled(Rule::InDictKeys) {
-                for generator in generators {
-                    flake8_simplify::rules::key_in_dict_for(
-                        checker,
-                        &generator.target,
-                        &generator.iter,
-                    );
-                }
-            }
             if checker.enabled(Rule::IterationOverSet) {
                 for generator in generators {
                     pylint::rules::iteration_over_set(checker, &generator.iter);

crates/ruff/src/checkers/ast/analyze/mod.rs

@@ -1,6 +1,7 @@
 pub(super) use argument::argument;
 pub(super) use arguments::arguments;
 pub(super) use bindings::bindings;
+pub(super) use comprehension::comprehension;
 pub(super) use deferred_for_loops::deferred_for_loops;
 pub(super) use deferred_scopes::deferred_scopes;
 pub(super) use definitions::definitions;
@@ -14,6 +15,7 @@ pub(super) use unresolved_references::unresolved_references;
 mod argument;
 mod arguments;
 mod bindings;
+mod comprehension;
 mod deferred_for_loops;
 mod deferred_scopes;
 mod definitions;

crates/ruff/src/checkers/ast/analyze/statement.rs

@@ -34,7 +34,6 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
                     pycodestyle::rules::ambiguous_variable_name(name, name.range())
                 }));
             }
-
             if checker.enabled(Rule::NonlocalWithoutBinding) {
                 if !checker.semantic.scope_id.is_global() {
                     for name in names {

crates/ruff/src/checkers/ast/mod.rs

@@ -1343,6 +1343,10 @@ impl<'a> Checker<'a> {
     /// Visit a list of [`Comprehension`] nodes, assumed to be the comprehensions that compose a
     /// generator expression, like a list or set comprehension.
     fn visit_generators(&mut self, generators: &'a [Comprehension]) {
+        for generator in generators {
+            analyze::comprehension(generator, self);
+        }
+
         let mut generators = generators.iter();
 
         let Some(generator) = generators.next() else {