Support variable keys in static dictionary key rule (#9411)

Closes https://github.com/astral-sh/ruff/issues/9410.

crates/ruff_linter/resources/test/fixtures/ruff/RUF011.py

@@ -1,13 +1,17 @@
 data = ["some", "Data"]
+constant = 5
 
-# Ok
+# OK
 {value: value.upper() for value in data}
 {value.lower(): value.upper() for value in data}
-{v: v*v for v in range(10)}
-{(0, "a", v): v*v for v in range(10)} # Tuple with variable
+{v: v * v for v in range(10)}
+{(0, "a", v): v * v for v in range(10)}  # Tuple with variable
+{constant: value.upper() for value in data for constant in data}
 
 # Errors
 {"key": value.upper() for value in data}
 {True: value.upper() for value in data}
 {0: value.upper() for value in data}
-{(1, "a"): value.upper() for value in data} # constant tuple
+{(1, "a"): value.upper() for value in data}  # Constant tuple
+{constant: value.upper() for value in data}
+{constant + constant: value.upper() for value in data}

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

@@ -1396,12 +1396,14 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
                 refurb::rules::reimplemented_starmap(checker, &comp.into());
             }
         }
-        Expr::DictComp(ast::ExprDictComp {
+        Expr::DictComp(
+            dict_comp @ ast::ExprDictComp {
                 key,
                 value,
                 generators,
                 range: _,
-        }) => {
+            },
+        ) => {
             if checker.enabled(Rule::UnnecessaryListIndexLookup) {
                 pylint::rules::unnecessary_list_index_lookup_comprehension(checker, expr);
             }
@@ -1422,7 +1424,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
                 }
             }
             if checker.enabled(Rule::StaticKeyDictComprehension) {
-                ruff::rules::static_key_dict_comprehension(checker, key);
+                ruff::rules::static_key_dict_comprehension(checker, dict_comp);
             }
         }
         Expr::GeneratorExp(

crates/ruff_linter/src/rules/ruff/rules/static_key_dict_comprehension.rs

@@ -1,16 +1,18 @@
-use ruff_python_ast::Expr;
+use rustc_hash::FxHashMap;
 
-use crate::fix::snippet::SourceCodeSnippet;
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
-use ruff_python_ast::helpers::is_constant;
+use ruff_python_ast::helpers::NameFinder;
+use ruff_python_ast::visitor::Visitor;
+use ruff_python_ast::{self as ast, Expr};
 use ruff_text_size::Ranged;
 
 use crate::checkers::ast::Checker;
+use crate::fix::snippet::SourceCodeSnippet;
 
 /// ## What it does
 /// Checks for dictionary comprehensions that use a static key, like a string
-/// literal.
+/// literal or a variable defined outside the comprehension.
 ///
 /// ## Why is this bad?
 /// Using a static key (like a string literal) in a dictionary comprehension
@@ -46,13 +48,40 @@ impl Violation for StaticKeyDictComprehension {
 }
 
 /// RUF011
-pub(crate) fn static_key_dict_comprehension(checker: &mut Checker, key: &Expr) {
-    if is_constant(key) {
+pub(crate) fn static_key_dict_comprehension(checker: &mut Checker, dict_comp: &ast::ExprDictComp) {
+    // Collect the bound names in the comprehension's generators.
+    let names = {
+        let mut visitor = NameFinder::default();
+        for generator in &dict_comp.generators {
+            visitor.visit_expr(&generator.target);
+        }
+        visitor.names
+    };
+
+    if is_constant(&dict_comp.key, &names) {
         checker.diagnostics.push(Diagnostic::new(
             StaticKeyDictComprehension {
-                key: SourceCodeSnippet::from_str(checker.locator().slice(key)),
+                key: SourceCodeSnippet::from_str(checker.locator().slice(dict_comp.key.as_ref())),
             },
-            key.range(),
+            dict_comp.key.range(),
         ));
     }
 }
+
+/// Returns `true` if the given expression is a constant in the context of the dictionary
+/// comprehension.
+fn is_constant(key: &Expr, names: &FxHashMap<&str, &ast::ExprName>) -> bool {
+    match key {
+        Expr::Tuple(ast::ExprTuple { elts, .. }) => elts.iter().all(|elt| is_constant(elt, names)),
+        Expr::Name(ast::ExprName { id, .. }) => !names.contains_key(id.as_str()),
+        Expr::BinOp(ast::ExprBinOp { left, right, .. }) => {
+            is_constant(left, names) && is_constant(right, names)
+        }
+        Expr::BoolOp(ast::ExprBoolOp { values, .. }) => {
+            values.iter().all(|value| is_constant(value, names))
+        }
+        Expr::UnaryOp(ast::ExprUnaryOp { operand, .. }) => is_constant(operand, names),
+        expr if expr.is_literal_expr() => true,
+        _ => false,
+    }
+}

crates/ruff_linter/src/rules/ruff/snapshots/ruff_linter__rules__ruff__tests__RUF011_RUF011.py.snap

@@ -1,40 +1,60 @@
 ---
 source: crates/ruff_linter/src/rules/ruff/mod.rs
 ---
-RUF011.py:10:2: RUF011 Dictionary comprehension uses static key: `"key"`
+RUF011.py:12:2: RUF011 Dictionary comprehension uses static key: `"key"`
    |
- 9 | # Errors
-10 | {"key": value.upper() for value in data}
+11 | # Errors
+12 | {"key": value.upper() for value in data}
    |  ^^^^^ RUF011
-11 | {True: value.upper() for value in data}
-12 | {0: value.upper() for value in data}
+13 | {True: value.upper() for value in data}
+14 | {0: value.upper() for value in data}
    |
 
-RUF011.py:11:2: RUF011 Dictionary comprehension uses static key: `True`
+RUF011.py:13:2: RUF011 Dictionary comprehension uses static key: `True`
    |
- 9 | # Errors
-10 | {"key": value.upper() for value in data}
-11 | {True: value.upper() for value in data}
+11 | # Errors
+12 | {"key": value.upper() for value in data}
+13 | {True: value.upper() for value in data}
    |  ^^^^ RUF011
-12 | {0: value.upper() for value in data}
-13 | {(1, "a"): value.upper() for value in data} # constant tuple
+14 | {0: value.upper() for value in data}
+15 | {(1, "a"): value.upper() for value in data}  # Constant tuple
    |
 
-RUF011.py:12:2: RUF011 Dictionary comprehension uses static key: `0`
+RUF011.py:14:2: RUF011 Dictionary comprehension uses static key: `0`
    |
-10 | {"key": value.upper() for value in data}
-11 | {True: value.upper() for value in data}
-12 | {0: value.upper() for value in data}
+12 | {"key": value.upper() for value in data}
+13 | {True: value.upper() for value in data}
+14 | {0: value.upper() for value in data}
    |  ^ RUF011
-13 | {(1, "a"): value.upper() for value in data} # constant tuple
+15 | {(1, "a"): value.upper() for value in data}  # Constant tuple
+16 | {constant: value.upper() for value in data}
    |
 
-RUF011.py:13:2: RUF011 Dictionary comprehension uses static key: `(1, "a")`
+RUF011.py:15:2: RUF011 Dictionary comprehension uses static key: `(1, "a")`
    |
-11 | {True: value.upper() for value in data}
-12 | {0: value.upper() for value in data}
-13 | {(1, "a"): value.upper() for value in data} # constant tuple
+13 | {True: value.upper() for value in data}
+14 | {0: value.upper() for value in data}
+15 | {(1, "a"): value.upper() for value in data}  # Constant tuple
    |  ^^^^^^^^ RUF011
+16 | {constant: value.upper() for value in data}
+17 | {constant + constant: value.upper() for value in data}
+   |
+
+RUF011.py:16:2: RUF011 Dictionary comprehension uses static key: `constant`
+   |
+14 | {0: value.upper() for value in data}
+15 | {(1, "a"): value.upper() for value in data}  # Constant tuple
+16 | {constant: value.upper() for value in data}
+   |  ^^^^^^^^ RUF011
+17 | {constant + constant: value.upper() for value in data}
+   |
+
+RUF011.py:17:2: RUF011 Dictionary comprehension uses static key: `constant + constant`
+   |
+15 | {(1, "a"): value.upper() for value in data}  # Constant tuple
+16 | {constant: value.upper() for value in data}
+17 | {constant + constant: value.upper() for value in data}
+   |  ^^^^^^^^^^^^^^^^^^^ RUF011
    |
 
 

crates/ruff_python_ast/src/helpers.rs

@@ -1,6 +1,7 @@
 use std::borrow::Cow;
 use std::path::Path;
 
+use rustc_hash::FxHashMap;
 use smallvec::SmallVec;
 
 use ruff_python_trivia::CommentRanges;
@@ -891,6 +892,25 @@ pub fn resolve_imported_module_path<'a>(
     Some(Cow::Owned(qualified_path))
 }
 
+/// A [`Visitor`] to collect all [`Expr::Name`] nodes in an AST.
+#[derive(Debug, Default)]
+pub struct NameFinder<'a> {
+    /// A map from identifier to defining expression.
+    pub names: FxHashMap<&'a str, &'a ast::ExprName>,
+}
+
+impl<'a, 'b> Visitor<'b> for NameFinder<'a>
+where
+    'b: 'a,
+{
+    fn visit_expr(&mut self, expr: &'a Expr) {
+        if let Expr::Name(name) = expr {
+            self.names.insert(&name.id, name);
+        }
+        crate::visitor::walk_expr(self, expr);
+    }
+}
+
 /// A [`StatementVisitor`] that collects all `return` statements in a function or method.
 #[derive(Default)]
 pub struct ReturnStatementVisitor<'a> {