add generator/compehrension expressions

astral-sh · Oct 22, 2023 · a0cd6b1 · a0cd6b1
1 parent 221d3c1
commit a0cd6b1
Show file tree

Hide file tree

Showing 4 changed files with 241 additions and 105 deletions.
diff --git a/crates/ruff_linter/resources/test/fixtures/pylint/unnecessary_dict_index_lookup.py b/crates/ruff_linter/resources/test/fixtures/pylint/unnecessary_dict_index_lookup.py
@@ -1,6 +1,10 @@
 FRUITS = {"apple": 1, "orange": 10, "berry": 22}
 
 def fix_these():
+    [FRUITS[fruit_name] for fruit_name, fruit_count in FRUITS.items()]  # PLR1733
+    {FRUITS[fruit_name] for fruit_name, fruit_count in FRUITS.items()}  # PLR1733
+    {fruit_name: FRUITS[fruit_name] for fruit_name, fruit_count in FRUITS.items()}  # PLR1733
+
     for fruit_name, fruit_count in FRUITS.items():
         print(FRUITS[fruit_name])  # PLR1733
         blah = FRUITS[fruit_name]  # PLR1733
@@ -15,6 +19,10 @@ def dont_fix_these():
 
 
 def value_intentionally_unused():
+    [FRUITS[fruit_name] for fruit_name, _ in FRUITS.items()]  # PLR1733
+    {FRUITS[fruit_name] for fruit_name, _ in FRUITS.items()}  # PLR1733
+    {fruit_name: FRUITS[fruit_name] for fruit_name, _ in FRUITS.items()}  # PLR1733
+
     for fruit_name, _ in FRUITS.items():
         print(FRUITS[fruit_name])  # Ok
         blah = FRUITS[fruit_name]  # Ok

diff --git a/crates/ruff_linter/src/checkers/ast/analyze/expression.rs b/crates/ruff_linter/src/checkers/ast/analyze/expression.rs
@@ -1295,6 +1295,9 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
                 range: _,
             },
         ) => {
+            if checker.enabled(Rule::UnnecessaryDictIndexLookup) {
+                pylint::rules::unnecessary_dict_index_lookup_comprehension(checker, expr);
+            }
             if checker.enabled(Rule::UnnecessaryComprehension) {
                 flake8_comprehensions::rules::unnecessary_list_set_comprehension(
                     checker, expr, elt, generators,
@@ -1319,6 +1322,9 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
                 range: _,
             },
         ) => {
+            if checker.enabled(Rule::UnnecessaryDictIndexLookup) {
+                pylint::rules::unnecessary_dict_index_lookup_comprehension(checker, expr);
+            }
             if checker.enabled(Rule::UnnecessaryComprehension) {
                 flake8_comprehensions::rules::unnecessary_list_set_comprehension(
                     checker, expr, elt, generators,
@@ -1342,6 +1348,9 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
             generators,
             range: _,
         }) => {
+            if checker.enabled(Rule::UnnecessaryDictIndexLookup) {
+                pylint::rules::unnecessary_dict_index_lookup_comprehension(checker, expr);
+            }
             if checker.enabled(Rule::UnnecessaryComprehension) {
                 flake8_comprehensions::rules::unnecessary_dict_comprehension(
                     checker, expr, key, value, generators,
@@ -1366,6 +1375,9 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
                 range: _,
             },
         ) => {
+            if checker.enabled(Rule::UnnecessaryDictIndexLookup) {
+                pylint::rules::unnecessary_dict_index_lookup_comprehension(checker, expr);
+            }
             if checker.enabled(Rule::FunctionUsesLoopVariable) {
                 flake8_bugbear::rules::function_uses_loop_variable(checker, &Node::Expr(expr));
             }

diff --git a/crates/ruff_linter/src/rules/pylint/rules/unnecessary_dict_index_lookup.rs b/crates/ruff_linter/src/rules/pylint/rules/unnecessary_dict_index_lookup.rs
@@ -103,62 +103,118 @@ impl<'a> Visitor<'_> for SubscriptVisitor<'a> {
 
 /// PLR1733
 pub(crate) fn unnecessary_dict_index_lookup(checker: &mut Checker, stmt_for: &StmtFor) {
+    let Some((dict_name, index_name, value_name)) = dict_items(&stmt_for.iter, &stmt_for.target)
+    else {
+        return;
+    };
+
+    let mut visitor = SubscriptVisitor::new(&dict_name, &index_name);
+
+    visitor.visit_body(&stmt_for.body);
+    visitor.visit_body(&stmt_for.orelse);
+
+    for range in visitor.diagnostic_ranges {
+        let mut diagnostic = Diagnostic::new(UnnecessaryDictIndexLookup, range);
+
+        diagnostic.set_fix(Fix::safe_edit(Edit::range_replacement(
+            value_name.clone(),
+            range,
+        )));
+
+        checker.diagnostics.push(diagnostic);
+    }
+}
+
+/// PLR1733
+pub(crate) fn unnecessary_dict_index_lookup_comprehension(checker: &mut Checker, expr: &Expr) {
+    match expr {
+        Expr::GeneratorExp(ast::ExprGeneratorExp {
+            elt, generators, ..
+        })
+        | Expr::DictComp(ast::ExprDictComp {
+            value: elt,
+            generators,
+            ..
+        })
+        | Expr::SetComp(ast::ExprSetComp {
+            elt, generators, ..
+        })
+        | Expr::ListComp(ast::ExprListComp {
+            elt, generators, ..
+        }) => {
+            for comp in generators {
+                let Some((dict_name, index_name, value_name)) =
+                    dict_items(&comp.iter, &comp.target)
+                else {
+                    continue;
+                };
+
+                let mut visitor = SubscriptVisitor::new(&dict_name, &index_name);
+
+                visitor.visit_expr(elt.as_ref());
+                for expr in &comp.ifs {
+                    visitor.visit_expr(expr);
+                }
+
+                for range in visitor.diagnostic_ranges {
+                    let mut diagnostic = Diagnostic::new(UnnecessaryDictIndexLookup, range);
+
+                    diagnostic.set_fix(Fix::safe_edit(Edit::range_replacement(
+                        value_name.clone(),
+                        range,
+                    )));
+
+                    checker.diagnostics.push(diagnostic);
+                }
+            }
+        }
+        _ => (),
+    }
+}
+
+fn dict_items(call_expr: &Expr, tuple_expr: &Expr) -> Option<(String, String, String)> {
     let Expr::Call(ast::ExprCall {
         func,
         arguments: Arguments { args, .. },
         ..
-    }) = stmt_for.iter.as_ref()
+    }) = call_expr
     else {
-        return;
+        return None;
     };
     if !args.is_empty() {
-        return;
+        return None;
     }
     let Expr::Attribute(ast::ExprAttribute { value, attr, .. }) = func.as_ref() else {
-        return;
+        return None;
     };
     if attr != "items" {
-        return;
+        return None;
     }
 
     let Expr::Name(ast::ExprName { id: dict_name, .. }) = value.as_ref() else {
-        return;
+        return None;
     };
 
-    let Expr::Tuple(ast::ExprTuple { elts, .. }) = stmt_for.target.as_ref() else {
-        return;
+    let Expr::Tuple(ast::ExprTuple { elts, .. }) = tuple_expr else {
+        return None;
     };
     let [index, value] = elts.as_slice() else {
-        return;
+        return None;
     };
 
     // Grab the variable names
     let Expr::Name(ast::ExprName { id: index_name, .. }) = index else {
-        return;
+        return None;
     };
 
     let Expr::Name(ast::ExprName { id: value_name, .. }) = value else {
-        return;
+        return None;
     };
 
     // If either of the variable names are intentionally ignored by naming them `_`, then don't emit
     if index_name == "_" || value_name == "_" {
-        return;
+        return None;
     }
 
-    let mut visitor = SubscriptVisitor::new(dict_name, index_name);
-
-    visitor.visit_body(&stmt_for.body);
-    visitor.visit_body(&stmt_for.orelse);
-
-    for range in visitor.diagnostic_ranges {
-        let mut diagnostic = Diagnostic::new(UnnecessaryDictIndexLookup, range);
-
-        diagnostic.set_fix(Fix::safe_edit(Edit::range_replacement(
-            value_name.clone(),
-            range,
-        )));
-
-        checker.diagnostics.push(diagnostic);
-    }
+    Some((dict_name.clone(), index_name.clone(), value_name.clone()))
 }