Skip to content

Commit

Permalink
Support variable keys in static dictionary key rule
Browse files Browse the repository at this point in the history
  • Loading branch information
charliermarsh committed Jan 6, 2024
1 parent c2c9997 commit 3836069
Show file tree
Hide file tree
Showing 5 changed files with 95 additions and 38 deletions.
11 changes: 7 additions & 4 deletions crates/ruff_linter/resources/test/fixtures/ruff/RUF011.py
Original file line number Diff line number Diff line change
@@ -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}
16 changes: 9 additions & 7 deletions crates/ruff_linter/src/checkers/ast/analyze/expression.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand All @@ -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(
Expand Down
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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,
}
}
Original file line number Diff line number Diff line change
@@ -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
|

Expand Down
20 changes: 20 additions & 0 deletions crates/ruff_python_ast/src/helpers.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
use std::borrow::Cow;
use std::path::Path;

use rustc_hash::FxHashMap;
use smallvec::SmallVec;

use ruff_python_trivia::CommentRanges;
Expand Down Expand Up @@ -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> {
Expand Down

0 comments on commit 3836069

Please sign in to comment.