Skip to content

Commit

Permalink
implement PLR6201 with autofix (#7973)
Browse files Browse the repository at this point in the history
## Summary

Implements
[R6201](https://pylint.readthedocs.io/en/latest/user_guide/messages/refactor/use-set-for-membership.html)
along with the autofix!

See #970 

## Test Plan

`cargo test`, and manually
  • Loading branch information
diceroll123 committed Oct 17, 2023
1 parent cb6d74c commit 5da0f91
Show file tree
Hide file tree
Showing 8 changed files with 154 additions and 0 deletions.
@@ -0,0 +1,10 @@
# Errors
1 in [1, 2, 3]
1 in (1, 2, 3)
1 in (
1, 2, 3
)

# OK
fruits = ["cherry", "grapes"]
"cherry" in fruits
3 changes: 3 additions & 0 deletions crates/ruff_linter/src/checkers/ast/analyze/expression.rs
Expand Up @@ -1197,6 +1197,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::LiteralMembership) {
pylint::rules::literal_membership(checker, compare);
}
if checker.enabled(Rule::ComparisonOfConstant) {
pylint::rules::comparison_of_constant(checker, left, ops, comparators);
}
Expand Down
1 change: 1 addition & 0 deletions crates/ruff_linter/src/codes.rs
Expand Up @@ -254,6 +254,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> {
(Pylint, "R1722") => (RuleGroup::Stable, rules::pylint::rules::SysExitAlias),
(Pylint, "R2004") => (RuleGroup::Stable, rules::pylint::rules::MagicValueComparison),
(Pylint, "R5501") => (RuleGroup::Stable, rules::pylint::rules::CollapsibleElseIf),
(Pylint, "R6201") => (RuleGroup::Preview, rules::pylint::rules::LiteralMembership),
#[allow(deprecated)]
(Pylint, "R6301") => (RuleGroup::Nursery, rules::pylint::rules::NoSelfUse),
(Pylint, "W0120") => (RuleGroup::Stable, rules::pylint::rules::UselessElseOnLoop),
Expand Down
1 change: 1 addition & 0 deletions crates/ruff_linter/src/rules/pylint/mod.rs
Expand Up @@ -137,6 +137,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::MisplacedBareRaise, Path::new("misplaced_bare_raise.py"))]
#[test_case(Rule::LiteralMembership, Path::new("literal_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(
Expand Down
65 changes: 65 additions & 0 deletions crates/ruff_linter/src/rules/pylint/rules/literal_membership.rs
@@ -0,0 +1,65 @@
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;

use crate::checkers::ast::Checker;

/// ## What it does
/// Checks for membership tests on `list` and `tuple` literals.
///
/// ## Why is this bad?
/// When testing for membership in a static sequence, prefer a `set` literal
/// over a `list` or `tuple`, as Python optimizes `set` membership tests.
///
/// ## Example
/// ```python
/// 1 in [1, 2, 3]
/// ```
///
/// Use instead:
/// ```python
/// 1 in {1, 2, 3}
/// ```
/// ## References
/// - [What’s New In Python 3.2](https://docs.python.org/3/whatsnew/3.2.html#optimizations)
#[violation]
pub struct LiteralMembership;

impl AlwaysFixableViolation for LiteralMembership {
#[derive_message_formats]
fn message(&self) -> String {
format!("Use a `set` literal when testing for membership")
}

fn fix_title(&self) -> String {
format!("Convert to `set`")
}
}

/// PLR6201
pub(crate) fn literal_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;
};

if !matches!(right, Expr::List(_) | Expr::Tuple(_)) {
return;
}

let mut diagnostic = Diagnostic::new(LiteralMembership, right.range());

let literal = checker.locator().slice(right);
let set = format!("{{{}}}", &literal[1..literal.len() - 1]);
diagnostic.set_fix(Fix::safe_edit(Edit::range_replacement(set, right.range())));

checker.diagnostics.push(diagnostic);
}
2 changes: 2 additions & 0 deletions crates/ruff_linter/src/rules/pylint/rules/mod.rs
Expand Up @@ -24,6 +24,7 @@ pub(crate) use invalid_envvar_value::*;
pub(crate) use invalid_str_return::*;
pub(crate) use invalid_string_characters::*;
pub(crate) use iteration_over_set::*;
pub(crate) use literal_membership::*;
pub(crate) use load_before_global_declaration::*;
pub(crate) use logging::*;
pub(crate) use magic_value_comparison::*;
Expand Down Expand Up @@ -86,6 +87,7 @@ mod invalid_envvar_value;
mod invalid_str_return;
mod invalid_string_characters;
mod iteration_over_set;
mod literal_membership;
mod load_before_global_declaration;
mod logging;
mod magic_value_comparison;
Expand Down
@@ -0,0 +1,69 @@
---
source: crates/ruff_linter/src/rules/pylint/mod.rs
---
literal_membership.py:2:6: PLR6201 [*] Use a `set` literal when testing for membership
|
1 | # Errors
2 | 1 in [1, 2, 3]
| ^^^^^^^^^ PLR6201
3 | 1 in (1, 2, 3)
4 | 1 in (
|
= help: Convert to `set`

Fix
1 1 | # Errors
2 |-1 in [1, 2, 3]
2 |+1 in {1, 2, 3}
3 3 | 1 in (1, 2, 3)
4 4 | 1 in (
5 5 | 1, 2, 3

literal_membership.py:3:6: PLR6201 [*] Use a `set` literal when testing for membership
|
1 | # Errors
2 | 1 in [1, 2, 3]
3 | 1 in (1, 2, 3)
| ^^^^^^^^^ PLR6201
4 | 1 in (
5 | 1, 2, 3
|
= help: Convert to `set`

Fix
1 1 | # Errors
2 2 | 1 in [1, 2, 3]
3 |-1 in (1, 2, 3)
3 |+1 in {1, 2, 3}
4 4 | 1 in (
5 5 | 1, 2, 3
6 6 | )

literal_membership.py:4:6: PLR6201 [*] Use a `set` literal when testing for membership
|
2 | 1 in [1, 2, 3]
3 | 1 in (1, 2, 3)
4 | 1 in (
| ______^
5 | | 1, 2, 3
6 | | )
| |_^ PLR6201
7 |
8 | # OK
|
= help: Convert to `set`

Fix
1 1 | # Errors
2 2 | 1 in [1, 2, 3]
3 3 | 1 in (1, 2, 3)
4 |-1 in (
4 |+1 in {
5 5 | 1, 2, 3
6 |-)
6 |+}
7 7 |
8 8 | # OK
9 9 | fruits = ["cherry", "grapes"]


3 changes: 3 additions & 0 deletions ruff.schema.json

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

0 comments on commit 5da0f91

Please sign in to comment.