[perflint] Extend PERF102 to comprehensions and generators

crates/ruff_linter/resources/test/fixtures/perflint/PERF102.py

@@ -105,3 +105,18 @@ def f():
 def _create_context(name_to_value):
     for(B,D)in A.items():
         if(C:=name_to_value.get(B.name)):A.run(B.set,C)
+
+
+# Comprehensions and generators — errors (https://github.com/astral-sh/ruff/issues/6638)
+_ = [k for k, _ in some_dict.items()]  # PERF102
+_ = {k for k, _ in some_dict.items()}  # PERF102
+_ = {k: "v" for k, _ in some_dict.items()}  # PERF102
+_ = (k for k, _ in some_dict.items())  # PERF102
+_ = [v for _, v in some_dict.items()]  # PERF102
+_ = [k for k, v in some_dict.items()]  # PERF102 (v unused)
+_ = [v for x in range(1) for _, v in some_dict.items()]  # PERF102
+
+# Comprehensions — no errors
+_ = [(k, v) for k, v in some_dict.items()]  # OK (both used)
+_ = [item for item in some_dict.items()]  # OK (not tuple target)
+_ = [k for k, v in some_dict.items() if v]  # OK (v used in condition)

crates/ruff_linter/src/checkers/ast/analyze/deferred_comprehensions.rs

@@ -0,0 +1,37 @@
+use ruff_python_ast::Expr;
+
+use crate::checkers::ast::Checker;
+use crate::codes::Rule;
+use crate::rules::perflint;
+
+/// Run lint rules over all deferred comprehensions in the [`SemanticModel`].
+pub(crate) fn deferred_comprehensions(checker: &mut Checker) {
+    while !checker.analyze.comprehensions.is_empty() {
+        let comprehensions = std::mem::take(&mut checker.analyze.comprehensions);
+        for snapshot in comprehensions {
+            checker.semantic.restore(snapshot);
+
+            let Some(generators) =
+                checker
+                    .semantic
+                    .current_expression()
+                    .and_then(|expr| match expr {
+                        Expr::ListComp(comp) => Some(comp.generators.as_slice()),
+                        Expr::SetComp(comp) => Some(comp.generators.as_slice()),
+                        Expr::DictComp(comp) => Some(comp.generators.as_slice()),
+                        Expr::Generator(generator) => Some(generator.generators.as_slice()),
+                        _ => None,
+                    })
+            else {
+                debug_assert!(false, "Expected a comprehension");
+                continue;
+            };
+
+            for generator in generators {
+                if checker.is_rule_enabled(Rule::IncorrectDictIterator) {
+                    perflint::rules::incorrect_dict_iterator_comprehension(checker, generator);
+                }
+            }
+        }
+    }
+}

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

@@ -1,5 +1,6 @@
 pub(super) use bindings::bindings;
 pub(super) use comprehension::comprehension;
+pub(super) use deferred_comprehensions::deferred_comprehensions;
 pub(super) use deferred_for_loops::deferred_for_loops;
 pub(super) use deferred_lambdas::deferred_lambdas;
 pub(super) use deferred_scopes::deferred_scopes;
@@ -16,6 +17,7 @@ pub(super) use unresolved_references::unresolved_references;
 
 mod bindings;
 mod comprehension;
+mod deferred_comprehensions;
 mod deferred_for_loops;
 mod deferred_lambdas;
 mod deferred_scopes;

crates/ruff_linter/src/checkers/ast/deferred.rs

@@ -34,4 +34,5 @@ pub(crate) struct Analyze {
     pub(crate) scopes: Vec<ScopeId>,
     pub(crate) lambdas: Vec<Snapshot>,
     pub(crate) for_loops: Vec<Snapshot>,
+    pub(crate) comprehensions: Vec<Snapshot>,
 }

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

@@ -68,7 +68,9 @@ use crate::docstrings::extraction::ExtractionTarget;
 use crate::importer::{ImportRequest, Importer, ResolutionError};
 use crate::noqa::NoqaMapping;
 use crate::package::PackageRoot;
-use crate::preview::is_undefined_export_in_dunder_init_enabled;
+use crate::preview::{
+    is_incorrect_dict_iterator_comprehension_enabled, is_undefined_export_in_dunder_init_enabled,
+};
 use crate::registry::Rule;
 use crate::rules::flake8_bugbear::rules::ReturnInGenerator;
 use crate::rules::pyflakes::rules::{
@@ -2465,6 +2467,12 @@ impl<'a> Checker<'a> {
         for generator in generators {
             analyze::comprehension(generator, self);
         }
+
+        if self.is_rule_enabled(Rule::IncorrectDictIterator)
+            && is_incorrect_dict_iterator_comprehension_enabled(self.settings())
+        {
+            self.analyze.comprehensions.push(self.semantic.snapshot());
+        }
     }
 
     /// Visit a body of [`Stmt`] nodes within a type-checking block.
@@ -3309,6 +3317,7 @@ pub(crate) fn check_ast(
     // Check docstrings, bindings, and unresolved references.
     analyze::deferred_lambdas(&mut checker);
     analyze::deferred_for_loops(&mut checker);
+    analyze::deferred_comprehensions(&mut checker);
     analyze::definitions(&mut checker);
     analyze::bindings(&checker);
     analyze::unresolved_references(&checker);

crates/ruff_linter/src/preview.rs

@@ -307,3 +307,10 @@ pub(crate) const fn is_expanded_import_conventions_enabled(preview: PreviewMode)
 pub(crate) const fn is_file_level_invalid_rule_code_enabled(settings: &LinterSettings) -> bool {
     settings.preview.is_enabled()
 }
+
+// https://github.com/astral-sh/ruff/pull/23473
+pub(crate) const fn is_incorrect_dict_iterator_comprehension_enabled(
+    settings: &LinterSettings,
+) -> bool {
+    settings.preview.is_enabled()
+}

crates/ruff_linter/src/rules/perflint/mod.rs

@@ -31,6 +31,7 @@ mod tests {
         Ok(())
     }
 
+    #[test_case(Rule::IncorrectDictIterator, Path::new("PERF102.py"))]
     // TODO: remove this test case when the fixes for `perf401` and `perf403` are stabilized
     #[test_case(Rule::ManualDictComprehension, Path::new("PERF403.py"))]
     #[test_case(Rule::ManualListComprehension, Path::new("PERF401.py"))]

crates/ruff_linter/src/rules/perflint/rules/incorrect_dict_iterator.rs

@@ -62,9 +62,21 @@ impl AlwaysFixableViolation for IncorrectDictIterator {
     }
 }
 
-/// PERF102
+/// PERF102 for `for` loops.
 pub(crate) fn incorrect_dict_iterator(checker: &Checker, stmt_for: &ast::StmtFor) {
-    let Expr::Tuple(ast::ExprTuple { elts, .. }) = stmt_for.target.as_ref() else {
+    check_dict_items_usage(checker, stmt_for.target.as_ref(), stmt_for.iter.as_ref());
+}
+
+/// PERF102 for comprehensions and generators.
+pub(crate) fn incorrect_dict_iterator_comprehension(
+    checker: &Checker,
+    comprehension: &ast::Comprehension,
+) {
+    check_dict_items_usage(checker, &comprehension.target, &comprehension.iter);
+}
+
+fn check_dict_items_usage(checker: &Checker, target: &Expr, iter: &Expr) {
+    let Expr::Tuple(ast::ExprTuple { elts, .. }) = target else {
         return;
     };
     let [key, value] = elts.as_slice() else {
@@ -74,7 +86,7 @@ pub(crate) fn incorrect_dict_iterator(checker: &Checker, stmt_for: &ast::StmtFor
         func,
         arguments: Arguments { args, .. },
         ..
-    }) = stmt_for.iter.as_ref()
+    }) = iter
     else {
         return;
     };
@@ -110,10 +122,10 @@ pub(crate) fn incorrect_dict_iterator(checker: &Checker, stmt_for: &ast::StmtFor
             let replace_target = Edit::range_replacement(
                 pad(
                     checker.locator().slice(value).to_string(),
-                    stmt_for.target.range(),
+                    target.range(),
                     checker.locator(),
                 ),
-                stmt_for.target.range(),
+                target.range(),
             );
             diagnostic.set_fix(Fix::unsafe_edits(replace_attribute, [replace_target]));
         }
@@ -129,10 +141,10 @@ pub(crate) fn incorrect_dict_iterator(checker: &Checker, stmt_for: &ast::StmtFor
             let replace_target = Edit::range_replacement(
                 pad(
                     checker.locator().slice(key).to_string(),
-                    stmt_for.target.range(),
+                    target.range(),
                     checker.locator(),
                 ),
-                stmt_for.target.range(),
+                target.range(),
             );
             diagnostic.set_fix(Fix::unsafe_edits(replace_attribute, [replace_target]));
         }

crates/ruff_linter/src/rules/perflint/snapshots/ruff_linter__rules__perflint__tests__PERF102_PERF102.py.snap

@@ -226,4 +226,6 @@ help: Replace `.items()` with `.keys()`
     -     for(B,D)in A.items():
 106 +     for B in A.keys():
 107 |         if(C:=name_to_value.get(B.name)):A.run(B.set,C)
+108 | 
+109 | 
 note: This is an unsafe fix and may change runtime behavior