From 3836069aca98bcce0384398498aa331b8329fbcf Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Sat, 6 Jan 2024 15:35:09 -0500 Subject: [PATCH] Support variable keys in static dictionary key rule --- .../resources/test/fixtures/ruff/RUF011.py | 11 ++-- .../src/checkers/ast/analyze/expression.rs | 16 +++--- .../rules/static_key_dict_comprehension.rs | 36 ++++++++++--- ..._rules__ruff__tests__RUF011_RUF011.py.snap | 50 +++++++++++-------- crates/ruff_python_ast/src/helpers.rs | 20 ++++++++ 5 files changed, 95 insertions(+), 38 deletions(-) diff --git a/crates/ruff_linter/resources/test/fixtures/ruff/RUF011.py b/crates/ruff_linter/resources/test/fixtures/ruff/RUF011.py index 3dce62bb4f476f..d94a4adb1e58e7 100644 --- a/crates/ruff_linter/resources/test/fixtures/ruff/RUF011.py +++ b/crates/ruff_linter/resources/test/fixtures/ruff/RUF011.py @@ -1,13 +1,16 @@ 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} diff --git a/crates/ruff_linter/src/checkers/ast/analyze/expression.rs b/crates/ruff_linter/src/checkers/ast/analyze/expression.rs index 4e9e75c45424d3..3b4286fd8a2906 100644 --- a/crates/ruff_linter/src/checkers/ast/analyze/expression.rs +++ b/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 { - key, - value, - generators, - range: _, - }) => { + 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( diff --git a/crates/ruff_linter/src/rules/ruff/rules/static_key_dict_comprehension.rs b/crates/ruff_linter/src/rules/ruff/rules/static_key_dict_comprehension.rs index 460495e262ea64..672e95c9530355 100644 --- a/crates/ruff_linter/src/rules/ruff/rules/static_key_dict_comprehension.rs +++ b/crates/ruff_linter/src/rules/ruff/rules/static_key_dict_comprehension.rs @@ -1,12 +1,14 @@ -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 @@ -46,13 +48,33 @@ 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 if expr.is_literal_expr() => true, + _ => false, + } +} diff --git a/crates/ruff_linter/src/rules/ruff/snapshots/ruff_linter__rules__ruff__tests__RUF011_RUF011.py.snap b/crates/ruff_linter/src/rules/ruff/snapshots/ruff_linter__rules__ruff__tests__RUF011_RUF011.py.snap index e522dfb7239860..30c213a62ae18c 100644 --- a/crates/ruff_linter/src/rules/ruff/snapshots/ruff_linter__rules__ruff__tests__RUF011_RUF011.py.snap +++ b/crates/ruff_linter/src/rules/ruff/snapshots/ruff_linter__rules__ruff__tests__RUF011_RUF011.py.snap @@ -1,39 +1,49 @@ --- source: crates/ruff_linter/src/rules/ruff/mod.rs --- -RUF011.py:10:2: RUF011 Dictionary comprehension uses static key: `"key"` +RUF011.py:11:2: RUF011 Dictionary comprehension uses static key: `"key"` | - 9 | # Errors -10 | {"key": value.upper() for value in data} +10 | # Errors +11 | {"key": value.upper() for value in data} | ^^^^^ RUF011 -11 | {True: value.upper() for value in data} -12 | {0: value.upper() for value in data} +12 | {True: value.upper() for value in data} +13 | {0: value.upper() for value in data} | -RUF011.py:11:2: RUF011 Dictionary comprehension uses static key: `True` +RUF011.py:12: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} +10 | # Errors +11 | {"key": value.upper() for value in data} +12 | {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 +13 | {0: value.upper() for value in data} +14 | {(1, "a"): value.upper() for value in data} # Constant tuple | -RUF011.py:12:2: RUF011 Dictionary comprehension uses static key: `0` +RUF011.py:13: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} +11 | {"key": value.upper() for value in data} +12 | {True: value.upper() for value in data} +13 | {0: value.upper() for value in data} | ^ RUF011 -13 | {(1, "a"): value.upper() for value in data} # constant tuple +14 | {(1, "a"): value.upper() for value in data} # Constant tuple +15 | {constant: value.upper() for value in data} | -RUF011.py:13:2: RUF011 Dictionary comprehension uses static key: `(1, "a")` +RUF011.py:14: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 +12 | {True: value.upper() for value in data} +13 | {0: value.upper() for value in data} +14 | {(1, "a"): value.upper() for value in data} # Constant tuple + | ^^^^^^^^ RUF011 +15 | {constant: value.upper() for value in data} + | + +RUF011.py:15:2: RUF011 Dictionary comprehension uses static key: `constant` + | +13 | {0: value.upper() for value in data} +14 | {(1, "a"): value.upper() for value in data} # Constant tuple +15 | {constant: value.upper() for value in data} | ^^^^^^^^ RUF011 | diff --git a/crates/ruff_python_ast/src/helpers.rs b/crates/ruff_python_ast/src/helpers.rs index 4efe25e3469a23..bb463c95b3f451 100644 --- a/crates/ruff_python_ast/src/helpers.rs +++ b/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> {