implement PLR6201 with autofix

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

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);
             }

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),

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(

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;

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);
+}

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"]
+
+