[pylint] Implement unnecessary-dict-index-lookup (PLR1733)

crates/ruff_linter/resources/test/fixtures/pylint/unnecessary_dict_index_lookup.py

@@ -0,0 +1,29 @@
+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
+        assert FRUITS[fruit_name] == "pear"  # PLR1733
+
+
+def dont_fix_these():
+    # once there is an assignment to the dict[index], we stop emitting diagnostics
+    for fruit_name, fruit_count in FRUITS.items():
+        FRUITS[fruit_name] = 0  # Ok
+        assert FRUITS[fruit_name] == 0  # Ok
+
+
+def value_intentionally_unused():
+    [FRUITS[fruit_name] for fruit_name, _ in FRUITS.items()]  # Ok
+    {FRUITS[fruit_name] for fruit_name, _ in FRUITS.items()}  # Ok
+    {fruit_name: FRUITS[fruit_name] for fruit_name, _ in FRUITS.items()}  # Ok
+
+    for fruit_name, _ in FRUITS.items():
+        print(FRUITS[fruit_name])  # Ok
+        blah = FRUITS[fruit_name]  # Ok
+        assert FRUITS[fruit_name] == "pear"  # Ok

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

@@ -1333,6 +1333,9 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
             if checker.enabled(Rule::UnnecessaryListIndexLookup) {
                 pylint::rules::unnecessary_list_index_lookup_comprehension(checker, expr);
             }
+            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,
@@ -1360,6 +1363,9 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
             if checker.enabled(Rule::UnnecessaryListIndexLookup) {
                 pylint::rules::unnecessary_list_index_lookup_comprehension(checker, expr);
             }
+            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,
@@ -1386,6 +1392,9 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
             if checker.enabled(Rule::UnnecessaryListIndexLookup) {
                 pylint::rules::unnecessary_list_index_lookup_comprehension(checker, expr);
             }
+            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,
@@ -1413,6 +1422,9 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
             if checker.enabled(Rule::UnnecessaryListIndexLookup) {
                 pylint::rules::unnecessary_list_index_lookup_comprehension(checker, expr);
             }
+            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));
             }

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

@@ -1280,6 +1280,9 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
             if checker.enabled(Rule::UnnecessaryListIndexLookup) {
                 pylint::rules::unnecessary_list_index_lookup(checker, for_stmt);
             }
+            if checker.enabled(Rule::UnnecessaryDictIndexLookup) {
+                pylint::rules::unnecessary_dict_index_lookup(checker, for_stmt);
+            }
             if !is_async {
                 if checker.enabled(Rule::ReimplementedBuiltin) {
                     flake8_simplify::rules::convert_for_loop_to_any_all(checker, stmt);

crates/ruff_linter/src/codes.rs

@@ -260,6 +260,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> {
         (Pylint, "R1714") => (RuleGroup::Stable, rules::pylint::rules::RepeatedEqualityComparison),
         (Pylint, "R1706") => (RuleGroup::Preview, rules::pylint::rules::AndOrTernary),
         (Pylint, "R1722") => (RuleGroup::Stable, rules::pylint::rules::SysExitAlias),
+        (Pylint, "R1733") => (RuleGroup::Preview, rules::pylint::rules::UnnecessaryDictIndexLookup),
         (Pylint, "R1736") => (RuleGroup::Preview, rules::pylint::rules::UnnecessaryListIndexLookup),
         (Pylint, "R2004") => (RuleGroup::Stable, rules::pylint::rules::MagicValueComparison),
         (Pylint, "R5501") => (RuleGroup::Stable, rules::pylint::rules::CollapsibleElseIf),

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

@@ -160,6 +160,10 @@ mod tests {
     )]
     #[test_case(Rule::NoClassmethodDecorator, Path::new("no_method_decorator.py"))]
     #[test_case(Rule::NoStaticmethodDecorator, Path::new("no_method_decorator.py"))]
+    #[test_case(
+        Rule::UnnecessaryDictIndexLookup,
+        Path::new("unnecessary_dict_index_lookup.py")
+    )]
     fn rules(rule_code: Rule, path: &Path) -> Result<()> {
         let snapshot = format!("{}_{}", rule_code.noqa_code(), path.to_string_lossy());
         let diagnostics = test_path(

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

@@ -62,6 +62,7 @@ pub(crate) use type_bivariance::*;
 pub(crate) use type_name_incorrect_variance::*;
 pub(crate) use type_param_name_mismatch::*;
 pub(crate) use unexpected_special_method_signature::*;
+pub(crate) use unnecessary_dict_index_lookup::*;
 pub(crate) use unnecessary_direct_lambda_call::*;
 pub(crate) use unnecessary_lambda::*;
 pub(crate) use unnecessary_list_index_lookup::*;
@@ -137,6 +138,7 @@ mod type_bivariance;
 mod type_name_incorrect_variance;
 mod type_param_name_mismatch;
 mod unexpected_special_method_signature;
+mod unnecessary_dict_index_lookup;
 mod unnecessary_direct_lambda_call;
 mod unnecessary_lambda;
 mod unnecessary_list_index_lookup;

crates/ruff_linter/src/rules/pylint/rules/unnecessary_dict_index_lookup.rs

@@ -0,0 +1,255 @@
+use ast::Stmt;
+use ruff_python_ast::{self as ast, Expr, StmtFor};
+
+use ruff_diagnostics::{AlwaysFixableViolation, Diagnostic, Edit, Fix};
+use ruff_macros::{derive_message_formats, violation};
+use ruff_python_ast::visitor;
+use ruff_python_ast::visitor::Visitor;
+use ruff_text_size::TextRange;
+
+use crate::checkers::ast::Checker;
+
+/// ## What it does
+/// Checks for index-based dict accesses during iterations through
+/// the `dict.items()`.
+///
+/// ## Why is this bad?
+/// It is more succinct to use the variable for the value at the current
+/// index which is already in scope from the iterator.
+///
+/// ## Example
+/// ```python
+/// FRUITS = {"apple": 1, "orange": 10, "berry": 22}
+///
+/// for fruit_name, fruit_count in FRUITS.items():
+///     print(FRUITS[fruit_name])
+/// ```
+///
+/// Use instead:
+/// ```python
+/// FRUITS = {"apple": 1, "orange": 10, "berry": 22}
+///
+/// for fruit_name, fruit_count in FRUITS.items():
+///     print(fruit_count)
+/// ```
+#[violation]
+pub struct UnnecessaryDictIndexLookup;
+
+impl AlwaysFixableViolation for UnnecessaryDictIndexLookup {
+    #[derive_message_formats]
+    fn message(&self) -> String {
+        format!("Unnecessary lookup of dict item by index")
+    }
+
+    fn fix_title(&self) -> String {
+        format!("Use existing variable")
+    }
+}
+
+/// 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.to_string(),
+            range,
+        )));
+
+        checker.diagnostics.push(diagnostic);
+    }
+}
+
+/// PLR1733
+pub(crate) fn unnecessary_dict_index_lookup_comprehension(checker: &mut Checker, expr: &Expr) {
+    let (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, ..
+    })) = expr
+    else {
+        return;
+    };
+
+    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.to_string(),
+                range,
+            )));
+
+            checker.diagnostics.push(diagnostic);
+        }
+    }
+}
+
+fn dict_items<'a>(
+    call_expr: &'a Expr,
+    tuple_expr: &'a Expr,
+) -> Option<(&'a str, &'a str, &'a str)> {
+    let ast::ExprCall {
+        func, arguments, ..
+    } = call_expr.as_call_expr()?;
+
+    if !arguments.is_empty() {
+        return None;
+    }
+    let Expr::Attribute(ast::ExprAttribute { value, attr, .. }) = func.as_ref() else {
+        return None;
+    };
+    if attr != "items" {
+        return None;
+    }
+
+    let Expr::Name(ast::ExprName { id: dict_name, .. }) = value.as_ref() else {
+        return None;
+    };
+
+    let Expr::Tuple(ast::ExprTuple { elts, .. }) = tuple_expr else {
+        return None;
+    };
+    let [index, value] = elts.as_slice() else {
+        return None;
+    };
+
+    // Grab the variable names
+    let Expr::Name(ast::ExprName { id: index_name, .. }) = index else {
+        return None;
+    };
+
+    let Expr::Name(ast::ExprName { id: value_name, .. }) = value else {
+        return None;
+    };
+
+    // If either of the variable names are intentionally ignored by naming them `_`, then don't emit
+    if index_name == "_" || value_name == "_" {
+        return None;
+    }
+
+    Some((dict_name, index_name, value_name))
+}
+
+#[derive(Debug)]
+struct SubscriptVisitor<'a> {
+    dict_name: &'a str,
+    index_name: &'a str,
+    diagnostic_ranges: Vec<TextRange>,
+    modified: bool,
+}
+
+impl<'a> SubscriptVisitor<'a> {
+    fn new(dict_name: &'a str, index_name: &'a str) -> Self {
+        Self {
+            dict_name,
+            index_name,
+            diagnostic_ranges: Vec::new(),
+            modified: false,
+        }
+    }
+}
+
+impl SubscriptVisitor<'_> {
+    fn check_target(&self, expr: &Expr) -> bool {
+        match expr {
+            Expr::Subscript(ast::ExprSubscript { value, slice, .. }) => {
+                let Expr::Name(ast::ExprName { id, .. }) = value.as_ref() else {
+                    return false;
+                };
+                if id == self.dict_name {
+                    let Expr::Name(ast::ExprName { id, .. }) = slice.as_ref() else {
+                        return false;
+                    };
+                    if id == self.index_name {
+                        return true;
+                    }
+                }
+                false
+            }
+            _ => false,
+        }
+    }
+}
+
+impl<'a> Visitor<'_> for SubscriptVisitor<'a> {
+    fn visit_stmt(&mut self, stmt: &Stmt) {
+        if self.modified {
+            return;
+        }
+        match stmt {
+            Stmt::Assign(ast::StmtAssign { targets, value, .. }) => {
+                self.modified = targets.iter().any(|target| self.check_target(target));
+                self.visit_expr(value);
+            }
+            Stmt::AnnAssign(ast::StmtAnnAssign { target, value, .. }) => {
+                if let Some(value) = value {
+                    self.modified = self.check_target(target);
+                    self.visit_expr(value);
+                }
+            }
+            Stmt::AugAssign(ast::StmtAugAssign { target, value, .. }) => {
+                self.modified = self.check_target(target);
+                self.visit_expr(value);
+            }
+            _ => visitor::walk_stmt(self, stmt),
+        }
+    }
+
+    fn visit_expr(&mut self, expr: &Expr) {
+        if self.modified {
+            return;
+        }
+        match expr {
+            Expr::Subscript(ast::ExprSubscript {
+                value,
+                slice,
+                range,
+                ..
+            }) => {
+                let Expr::Name(ast::ExprName { id, .. }) = value.as_ref() else {
+                    return;
+                };
+                if id == self.dict_name {
+                    let Expr::Name(ast::ExprName { id, .. }) = slice.as_ref() else {
+                        return;
+                    };
+                    if id == self.index_name {
+                        self.diagnostic_ranges.push(*range);
+                    }
+                }
+            }
+            _ => visitor::walk_expr(self, expr),
+        }
+    }
+}