From b07a1cbbd545830ef938ccb86382fe7cad724468 Mon Sep 17 00:00:00 2001 From: Steve C Date: Sun, 15 Oct 2023 23:49:56 -0400 Subject: [PATCH] implement PLR6201 with autofix --- .../test/fixtures/pylint/set_membership.py | 13 ++++ .../src/checkers/ast/analyze/expression.rs | 3 + crates/ruff_linter/src/codes.rs | 1 + crates/ruff_linter/src/rules/pylint/mod.rs | 1 + .../ruff_linter/src/rules/pylint/rules/mod.rs | 2 + .../src/rules/pylint/rules/set_membership.rs | 77 +++++++++++++++++++ ...int__tests__PLR6201_set_membership.py.snap | 63 +++++++++++++++ ruff.schema.json | 3 + 8 files changed, 163 insertions(+) create mode 100644 crates/ruff_linter/resources/test/fixtures/pylint/set_membership.py create mode 100644 crates/ruff_linter/src/rules/pylint/rules/set_membership.rs create mode 100644 crates/ruff_linter/src/rules/pylint/snapshots/ruff_linter__rules__pylint__tests__PLR6201_set_membership.py.snap diff --git a/crates/ruff_linter/resources/test/fixtures/pylint/set_membership.py b/crates/ruff_linter/resources/test/fixtures/pylint/set_membership.py new file mode 100644 index 00000000000000..9209558c64258e --- /dev/null +++ b/crates/ruff_linter/resources/test/fixtures/pylint/set_membership.py @@ -0,0 +1,13 @@ +# PLR6201 +1 in [1, 2, 3] + +# PLR6201 +1 in (1, 2, 3) + +# PLR6201 +def fruit_is_dangerous_for_cat(fruit: str) -> bool: + return fruit in ["cherry", "grapes"] + +# Ok +fruits = ["cherry", "grapes"] +"cherry" in fruits \ No newline at end of file diff --git a/crates/ruff_linter/src/checkers/ast/analyze/expression.rs b/crates/ruff_linter/src/checkers/ast/analyze/expression.rs index 58ec6f6e4bdcd2..dd70eb6c1ec695 100644 --- a/crates/ruff_linter/src/checkers/ast/analyze/expression.rs +++ b/crates/ruff_linter/src/checkers/ast/analyze/expression.rs @@ -1194,6 +1194,9 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) { if checker.enabled(Rule::ComparisonWithItself) { pylint::rules::comparison_with_itself(checker, left, ops, comparators); } + if checker.enabled(Rule::SetMembership) { + pylint::rules::set_membership(checker, compare); + } if checker.enabled(Rule::ComparisonOfConstant) { pylint::rules::comparison_of_constant(checker, left, ops, comparators); } diff --git a/crates/ruff_linter/src/codes.rs b/crates/ruff_linter/src/codes.rs index a0722c0741b9b0..cc92eab7c8ae10 100644 --- a/crates/ruff_linter/src/codes.rs +++ b/crates/ruff_linter/src/codes.rs @@ -239,6 +239,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> { (Pylint, "E2514") => (RuleGroup::Unspecified, rules::pylint::rules::InvalidCharacterNul), (Pylint, "E2515") => (RuleGroup::Unspecified, rules::pylint::rules::InvalidCharacterZeroWidthSpace), (Pylint, "R0124") => (RuleGroup::Unspecified, rules::pylint::rules::ComparisonWithItself), + (Pylint, "R6201") => (RuleGroup::Unspecified, rules::pylint::rules::SetMembership), (Pylint, "R0133") => (RuleGroup::Unspecified, rules::pylint::rules::ComparisonOfConstant), (Pylint, "R0206") => (RuleGroup::Unspecified, rules::pylint::rules::PropertyWithParameters), (Pylint, "R0402") => (RuleGroup::Unspecified, rules::pylint::rules::ManualFromImport), diff --git a/crates/ruff_linter/src/rules/pylint/mod.rs b/crates/ruff_linter/src/rules/pylint/mod.rs index 21b0c0250f911f..dca63cb6174cd0 100644 --- a/crates/ruff_linter/src/rules/pylint/mod.rs +++ b/crates/ruff_linter/src/rules/pylint/mod.rs @@ -135,6 +135,7 @@ mod tests { )] #[test_case(Rule::BadDunderMethodName, Path::new("bad_dunder_method_name.py"))] #[test_case(Rule::NoSelfUse, Path::new("no_self_use.py"))] + #[test_case(Rule::SetMembership, Path::new("set_membership.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_linter/src/rules/pylint/rules/mod.rs b/crates/ruff_linter/src/rules/pylint/rules/mod.rs index 317657266794bb..4f8ff68f40ce37 100644 --- a/crates/ruff_linter/src/rules/pylint/rules/mod.rs +++ b/crates/ruff_linter/src/rules/pylint/rules/mod.rs @@ -38,6 +38,7 @@ pub(crate) use repeated_equality_comparison::*; pub(crate) use repeated_isinstance_calls::*; pub(crate) use return_in_init::*; pub(crate) use self_assigning_variable::*; +pub(crate) use set_membership::*; pub(crate) use single_string_slots::*; pub(crate) use subprocess_popen_preexec_fn::*; pub(crate) use subprocess_run_without_check::*; @@ -98,6 +99,7 @@ mod repeated_equality_comparison; mod repeated_isinstance_calls; mod return_in_init; mod self_assigning_variable; +mod set_membership; mod single_string_slots; mod subprocess_popen_preexec_fn; mod subprocess_run_without_check; diff --git a/crates/ruff_linter/src/rules/pylint/rules/set_membership.rs b/crates/ruff_linter/src/rules/pylint/rules/set_membership.rs new file mode 100644 index 00000000000000..1951df9fdee0b2 --- /dev/null +++ b/crates/ruff_linter/src/rules/pylint/rules/set_membership.rs @@ -0,0 +1,77 @@ +use ruff_diagnostics::{AlwaysFixableViolation, Diagnostic, Edit, Fix}; +use ruff_macros::{derive_message_formats, violation}; +use ruff_python_ast::{self as ast, CmpOp, Expr}; +use ruff_text_size::{Ranged, TextRange}; + +use crate::checkers::ast::Checker; + +/// ## What it does +/// Checks for operations that checks a list or tuple for an element. +/// +/// ## Why is this bad? +/// Membership tests are more efficient when performed on a +/// lookup-optimized datatype like `set`. +/// +/// ## Example +/// ```python +/// 1 in [1, 2, 3] +/// ``` +/// +/// Use instead: +/// ```python +/// 1 in {1, 2, 3} +/// ``` +/// ## References +/// - [Python 3.2 release notes](https://docs.python.org/3/whatsnew/3.2.html#optimizations) +#[violation] +pub struct SetMembership; + +impl AlwaysFixableViolation for SetMembership { + #[derive_message_formats] + fn message(&self) -> String { + format!("Use a `set` when checking for element membership") + } + + fn fix_title(&self) -> String { + format!("Use a `set` when checking for element membership") + } +} + +/// PLR6201 +pub(crate) fn set_membership(checker: &mut Checker, compare: &ast::ExprCompare) { + let [op] = compare.ops.as_slice() else { + return; + }; + + if !matches!(op, CmpOp::In | CmpOp::NotIn) { + return; + } + + let [right] = compare.comparators.as_slice() else { + return; + }; + + let (Expr::List(ast::ExprList { + elts: right_elements, + .. + }) + | Expr::Tuple(ast::ExprTuple { + elts: right_elements, + .. + })) = right + else { + return; + }; + + let mut diagnostic = Diagnostic::new(SetMembership, right.range()); + + diagnostic.set_fix(Fix::safe_edit(Edit::range_replacement( + checker.generator().expr(&Expr::Set(ast::ExprSet { + elts: right_elements.clone(), + range: TextRange::default(), + })), + right.range(), + ))); + + checker.diagnostics.push(diagnostic); +} diff --git a/crates/ruff_linter/src/rules/pylint/snapshots/ruff_linter__rules__pylint__tests__PLR6201_set_membership.py.snap b/crates/ruff_linter/src/rules/pylint/snapshots/ruff_linter__rules__pylint__tests__PLR6201_set_membership.py.snap new file mode 100644 index 00000000000000..f7c702b5ba37d1 --- /dev/null +++ b/crates/ruff_linter/src/rules/pylint/snapshots/ruff_linter__rules__pylint__tests__PLR6201_set_membership.py.snap @@ -0,0 +1,63 @@ +--- +source: crates/ruff_linter/src/rules/pylint/mod.rs +--- +set_membership.py:2:6: PLR6201 [*] Use a `set` when checking for element membership + | +1 | # PLR6201 +2 | 1 in [1, 2, 3] + | ^^^^^^^^^ PLR6201 +3 | +4 | # PLR6201 + | + = help: Use a `set` when checking for element membership + +ℹ Fix +1 1 | # PLR6201 +2 |-1 in [1, 2, 3] + 2 |+1 in {1, 2, 3} +3 3 | +4 4 | # PLR6201 +5 5 | 1 in (1, 2, 3) + +set_membership.py:5:6: PLR6201 [*] Use a `set` when checking for element membership + | +4 | # PLR6201 +5 | 1 in (1, 2, 3) + | ^^^^^^^^^ PLR6201 +6 | +7 | # PLR6201 + | + = help: Use a `set` when checking for element membership + +ℹ Fix +2 2 | 1 in [1, 2, 3] +3 3 | +4 4 | # PLR6201 +5 |-1 in (1, 2, 3) + 5 |+1 in {1, 2, 3} +6 6 | +7 7 | # PLR6201 +8 8 | def fruit_is_dangerous_for_cat(fruit: str) -> bool: + +set_membership.py:9:21: PLR6201 [*] Use a `set` when checking for element membership + | + 7 | # PLR6201 + 8 | def fruit_is_dangerous_for_cat(fruit: str) -> bool: + 9 | return fruit in ["cherry", "grapes"] + | ^^^^^^^^^^^^^^^^^^^^ PLR6201 +10 | +11 | # Ok + | + = help: Use a `set` when checking for element membership + +ℹ Fix +6 6 | +7 7 | # PLR6201 +8 8 | def fruit_is_dangerous_for_cat(fruit: str) -> bool: +9 |- return fruit in ["cherry", "grapes"] + 9 |+ return fruit in {"cherry", "grapes"} +10 10 | +11 11 | # Ok +12 12 | fruits = ["cherry", "grapes"] + + diff --git a/ruff.schema.json b/ruff.schema.json index 7be275a4100e0f..1da4ad77073552 100644 --- a/ruff.schema.json +++ b/ruff.schema.json @@ -2975,6 +2975,9 @@ "PLR550", "PLR5501", "PLR6", + "PLR62", + "PLR620", + "PLR6201", "PLR63", "PLR630", "PLR6301",