diff --git a/crates/ruff/resources/test/fixtures/perflint/PERF403.py b/crates/ruff/resources/test/fixtures/perflint/PERF403.py new file mode 100644 index 0000000000..e28723673a --- /dev/null +++ b/crates/ruff/resources/test/fixtures/perflint/PERF403.py @@ -0,0 +1,83 @@ +def foo(): + fruit = ["apple", "pear", "orange"] + result = {} + for idx, name in enumerate(fruit): + result[idx] = name # PERF403 + + +def foo(): + fruit = ["apple", "pear", "orange"] + result = {} + for idx, name in enumerate(fruit): + if idx % 2: + result[idx] = name # PERF403 + + +def foo(): + fruit = ["apple", "pear", "orange"] + result = {} + for idx, name in enumerate(fruit): + if idx % 2: + result[idx] = name # Ok (false negative: edge case where `else` is same as `if`) + else: + result[idx] = name + + +def foo(): + result = {} + fruit = ["apple", "pear", "orange"] + for idx, name in enumerate(fruit): + if idx % 2: + result[idx] = name # PERF403 + + +def foo(): + fruit = ["apple", "pear", "orange"] + result = [] + for idx, name in enumerate(fruit): + if idx % 2: + result[idx] = name # OK (result is not a dictionary) + else: + result[idx] = name + + +def foo(): + fruit = ["apple", "pear", "orange"] + result = {} + for idx, name in enumerate(fruit): + if idx % 2: + result[idx] = name # OK (if/elif/else isn't replaceable) + elif idx % 3: + result[idx] = name + else: + result[idx] = name + + +def foo(): + result = {1: "banana"} + fruit = ["apple", "pear", "orange"] + for idx, name in enumerate(fruit): + if idx % 2: + result[idx] = name # PERF403 + + +def foo(): + fruit = ["apple", "pear", "orange"] + result = {} + for idx, name in enumerate(fruit): + if idx in result: + result[idx] = name # PERF403 + + +def foo(): + fruit = ["apple", "pear", "orange"] + result = {} + for name in fruit: + result[name] = name # PERF403 + + +def foo(): + fruit = ["apple", "pear", "orange"] + result = {} + for idx, name in enumerate(fruit): + result[name] = idx # PERF403 diff --git a/crates/ruff/src/checkers/ast/analyze/statement.rs b/crates/ruff/src/checkers/ast/analyze/statement.rs index f3c9596331..9b1bbddee4 100644 --- a/crates/ruff/src/checkers/ast/analyze/statement.rs +++ b/crates/ruff/src/checkers/ast/analyze/statement.rs @@ -1184,7 +1184,7 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) { iter, orelse, is_async, - .. + range: _, }, ) => { if checker.any_enabled(&[ @@ -1218,6 +1218,9 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) { if checker.enabled(Rule::ManualListCopy) { perflint::rules::manual_list_copy(checker, target, body); } + if checker.enabled(Rule::ManualDictComprehension) { + perflint::rules::manual_dict_comprehension(checker, target, body); + } if checker.enabled(Rule::UnnecessaryListCast) { perflint::rules::unnecessary_list_cast(checker, iter); } diff --git a/crates/ruff/src/codes.rs b/crates/ruff/src/codes.rs index c8c75ab2ef..45e8a63188 100644 --- a/crates/ruff/src/codes.rs +++ b/crates/ruff/src/codes.rs @@ -898,6 +898,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> { (Perflint, "203") => (RuleGroup::Unspecified, rules::perflint::rules::TryExceptInLoop), (Perflint, "401") => (RuleGroup::Unspecified, rules::perflint::rules::ManualListComprehension), (Perflint, "402") => (RuleGroup::Unspecified, rules::perflint::rules::ManualListCopy), + (Perflint, "403") => (RuleGroup::Unspecified, rules::perflint::rules::ManualDictComprehension), // flake8-fixme (Flake8Fixme, "001") => (RuleGroup::Unspecified, rules::flake8_fixme::rules::LineContainsFixme), diff --git a/crates/ruff/src/rules/perflint/mod.rs b/crates/ruff/src/rules/perflint/mod.rs index 3e196d479d..0f060aefa1 100644 --- a/crates/ruff/src/rules/perflint/mod.rs +++ b/crates/ruff/src/rules/perflint/mod.rs @@ -19,6 +19,7 @@ mod tests { #[test_case(Rule::TryExceptInLoop, Path::new("PERF203.py"))] #[test_case(Rule::ManualListComprehension, Path::new("PERF401.py"))] #[test_case(Rule::ManualListCopy, Path::new("PERF402.py"))] + #[test_case(Rule::ManualDictComprehension, Path::new("PERF403.py"))] fn rules(rule_code: Rule, path: &Path) -> Result<()> { let snapshot = format!("{}_{}", rule_code.noqa_code(), path.to_string_lossy()); let diagnostics = test_path( diff --git a/crates/ruff/src/rules/perflint/rules/manual_dict_comprehension.rs b/crates/ruff/src/rules/perflint/rules/manual_dict_comprehension.rs new file mode 100644 index 0000000000..559ba86e1e --- /dev/null +++ b/crates/ruff/src/rules/perflint/rules/manual_dict_comprehension.rs @@ -0,0 +1,178 @@ +use ruff_diagnostics::{Diagnostic, Violation}; +use ruff_macros::{derive_message_formats, violation}; +use ruff_python_ast::comparable::ComparableExpr; +use ruff_python_ast::helpers::any_over_expr; +use ruff_python_ast::{self as ast, Expr, Stmt}; +use ruff_python_semantic::analyze::typing::is_dict; +use ruff_python_semantic::Binding; + +use crate::checkers::ast::Checker; + +/// ## What it does +/// Checks for `for` loops that can be replaced by a dictionary comprehension. +/// +/// ## Why is this bad? +/// When creating or extending a dictionary in a for-loop, prefer a dictionary +/// comprehension. Comprehensions are more readable and more performant. +/// +/// For example, when comparing `{x: x for x in list(range(1000))}` to the `for` +/// loop version, the comprehension is ~10% faster on Python 3.11. +/// +/// Note that, as with all `perflint` rules, this is only intended as a +/// micro-optimization, and will have a negligible impact on performance in +/// most cases. +/// +/// ## Example +/// ```python +/// pairs = (("a", 1), ("b", 2)) +/// result = {} +/// for x, y in pairs: +/// if y % 2: +/// result[x] = y +/// ``` +/// +/// Use instead: +/// ```python +/// pairs = (("a", 1), ("b", 2)) +/// result = {x: y for x, y in pairs if y % 2} +/// ``` +/// +/// If you're appending to an existing dictionary, use the `update` method instead: +/// ```python +/// pairs = (("a", 1), ("b", 2)) +/// result.update({x: y for x, y in pairs if y % 2}) +/// ``` +#[violation] +pub struct ManualDictComprehension; + +impl Violation for ManualDictComprehension { + #[derive_message_formats] + fn message(&self) -> String { + format!("Use a dictionary comprehension instead of a for-loop") + } +} + +/// PERF403 +pub(crate) fn manual_dict_comprehension(checker: &mut Checker, target: &Expr, body: &[Stmt]) { + let (stmt, if_test) = match body { + // ```python + // for idx, name in enumerate(names): + // if idx % 2 == 0: + // result[name] = idx + // ``` + [Stmt::If(ast::StmtIf { + body, + elif_else_clauses, + test, + .. + })] => { + // TODO(charlie): If there's an `else` clause, verify that the `else` has the + // same structure. + if !elif_else_clauses.is_empty() { + return; + } + let [stmt] = body.as_slice() else { + return; + }; + (stmt, Some(test)) + } + // ```python + // for idx, name in enumerate(names): + // result[name] = idx + // ``` + [stmt] => (stmt, None), + _ => return, + }; + + let Stmt::Assign(ast::StmtAssign { + targets, + value, + range, + }) = stmt + else { + return; + }; + + let [Expr::Subscript(ast::ExprSubscript { + value: subscript_value, + slice, + .. + })] = targets.as_slice() + else { + return; + }; + + match target { + Expr::Tuple(ast::ExprTuple { elts, .. }) => { + if !elts + .iter() + .any(|elt| ComparableExpr::from(slice) == ComparableExpr::from(elt)) + { + return; + } + if !elts + .iter() + .any(|elt| ComparableExpr::from(value) == ComparableExpr::from(elt)) + { + return; + } + } + Expr::Name(_) => { + if ComparableExpr::from(slice) != ComparableExpr::from(target) { + return; + } + if ComparableExpr::from(value) != ComparableExpr::from(target) { + return; + } + } + _ => return, + } + + // Exclude non-dictionary value. + let Expr::Name(ast::ExprName { + id: subscript_name, .. + }) = subscript_value.as_ref() + else { + return; + }; + let scope = checker.semantic().current_scope(); + let bindings: Vec<&Binding> = scope + .get_all(subscript_name) + .map(|binding_id| checker.semantic().binding(binding_id)) + .collect(); + + let [binding] = bindings.as_slice() else { + return; + }; + + if !is_dict(binding, checker.semantic()) { + return; + } + + // Avoid if the value is used in the conditional test, e.g., + // + // ```python + // for x in y: + // if x in filtered: + // filtered[x] = y + // ``` + // + // Converting this to a dictionary comprehension would raise a `NameError` as + // `filtered` is not defined yet: + // + // ```python + // filtered = {x: y for x in y if x in filtered} + // ``` + if if_test.is_some_and(|test| { + any_over_expr(test, &|expr| { + expr.as_name_expr() + .is_some_and(|expr| expr.id == *subscript_name) + }) + }) { + return; + } + + checker + .diagnostics + .push(Diagnostic::new(ManualDictComprehension, *range)); +} diff --git a/crates/ruff/src/rules/perflint/rules/mod.rs b/crates/ruff/src/rules/perflint/rules/mod.rs index 690b0fc1fe..6da79326a4 100644 --- a/crates/ruff/src/rules/perflint/rules/mod.rs +++ b/crates/ruff/src/rules/perflint/rules/mod.rs @@ -1,10 +1,12 @@ pub(crate) use incorrect_dict_iterator::*; +pub(crate) use manual_dict_comprehension::*; pub(crate) use manual_list_comprehension::*; pub(crate) use manual_list_copy::*; pub(crate) use try_except_in_loop::*; pub(crate) use unnecessary_list_cast::*; mod incorrect_dict_iterator; +mod manual_dict_comprehension; mod manual_list_comprehension; mod manual_list_copy; mod try_except_in_loop; diff --git a/crates/ruff/src/rules/perflint/snapshots/ruff__rules__perflint__tests__PERF403_PERF403.py.snap b/crates/ruff/src/rules/perflint/snapshots/ruff__rules__perflint__tests__PERF403_PERF403.py.snap new file mode 100644 index 0000000000..461e20a438 --- /dev/null +++ b/crates/ruff/src/rules/perflint/snapshots/ruff__rules__perflint__tests__PERF403_PERF403.py.snap @@ -0,0 +1,52 @@ +--- +source: crates/ruff/src/rules/perflint/mod.rs +--- +PERF403.py:5:9: PERF403 Use a dictionary comprehension instead of a for-loop + | +3 | result = {} +4 | for idx, name in enumerate(fruit): +5 | result[idx] = name # PERF403 + | ^^^^^^^^^^^^^^^^^^ PERF403 + | + +PERF403.py:13:13: PERF403 Use a dictionary comprehension instead of a for-loop + | +11 | for idx, name in enumerate(fruit): +12 | if idx % 2: +13 | result[idx] = name # PERF403 + | ^^^^^^^^^^^^^^^^^^ PERF403 + | + +PERF403.py:31:13: PERF403 Use a dictionary comprehension instead of a for-loop + | +29 | for idx, name in enumerate(fruit): +30 | if idx % 2: +31 | result[idx] = name # PERF403 + | ^^^^^^^^^^^^^^^^^^ PERF403 + | + +PERF403.py:61:13: PERF403 Use a dictionary comprehension instead of a for-loop + | +59 | for idx, name in enumerate(fruit): +60 | if idx % 2: +61 | result[idx] = name # PERF403 + | ^^^^^^^^^^^^^^^^^^ PERF403 + | + +PERF403.py:76:9: PERF403 Use a dictionary comprehension instead of a for-loop + | +74 | result = {} +75 | for name in fruit: +76 | result[name] = name # PERF403 + | ^^^^^^^^^^^^^^^^^^^ PERF403 + | + +PERF403.py:83:9: PERF403 Use a dictionary comprehension instead of a for-loop + | +81 | result = {} +82 | for idx, name in enumerate(fruit): +83 | result[name] = idx # PERF403 + | ^^^^^^^^^^^^^^^^^^ PERF403 + | + + diff --git a/ruff.schema.json b/ruff.schema.json index 9223d40546..c672b5b592 100644 --- a/ruff.schema.json +++ b/ruff.schema.json @@ -2195,6 +2195,7 @@ "PERF40", "PERF401", "PERF402", + "PERF403", "PGH", "PGH0", "PGH00",