diff --git a/README.md b/README.md index c35850d499..2f493d0ee0 100644 --- a/README.md +++ b/README.md @@ -973,6 +973,7 @@ For more, see [flake8-simplify](https://pypi.org/project/flake8-simplify/0.19.3/ | SIM105 | UseContextlibSuppress | Use `contextlib.suppress(...)` instead of try-except-pass | | | SIM107 | ReturnInTryExceptFinally | Don't use `return` in `try`/`except` and `finally` | | | SIM108 | UseTernaryOperator | Use ternary operator `..` instead of if-else-block | 🛠 | +| SIM109 | CompareWithTuple | Use `value in (..., ...)` instead of `value == ... or value == ...` | 🛠 | | SIM110 | ConvertLoopToAny | Use `return any(x for x in y)` instead of `for` loop | 🛠 | | SIM111 | ConvertLoopToAll | Use `return all(x for x in y)` instead of `for` loop | 🛠 | | SIM117 | MultipleWithStatements | Use a single `with` statement with multiple contexts instead of nested `with` statements | | diff --git a/resources/test/fixtures/flake8_simplify/SIM109.py b/resources/test/fixtures/flake8_simplify/SIM109.py new file mode 100644 index 0000000000..c794e6fd84 --- /dev/null +++ b/resources/test/fixtures/flake8_simplify/SIM109.py @@ -0,0 +1,7 @@ +# Bad +if a == b or a == c: + d + +# Good +if a in (b, c): + d \ No newline at end of file diff --git a/ruff.schema.json b/ruff.schema.json index 0f69e17815..3fa37aa8e9 100644 --- a/ruff.schema.json +++ b/ruff.schema.json @@ -943,6 +943,7 @@ "SIM105", "SIM107", "SIM108", + "SIM109", "SIM11", "SIM110", "SIM111", diff --git a/src/checkers/ast.rs b/src/checkers/ast.rs index eb44221305..5f977d0071 100644 --- a/src/checkers/ast.rs +++ b/src/checkers/ast.rs @@ -2735,6 +2735,9 @@ where if self.settings.enabled.contains(&CheckCode::SIM101) { flake8_simplify::plugins::duplicate_isinstance_call(self, expr); } + if self.settings.enabled.contains(&CheckCode::SIM109) { + flake8_simplify::plugins::compare_with_tuple(self, expr); + } if self.settings.enabled.contains(&CheckCode::SIM220) { flake8_simplify::plugins::a_and_not_a(self, expr); } diff --git a/src/flake8_simplify/mod.rs b/src/flake8_simplify/mod.rs index 1966507bda..7b04ff0440 100644 --- a/src/flake8_simplify/mod.rs +++ b/src/flake8_simplify/mod.rs @@ -16,6 +16,7 @@ mod tests { #[test_case(CheckCode::SIM102, Path::new("SIM102.py"); "SIM102")] #[test_case(CheckCode::SIM105, Path::new("SIM105.py"); "SIM105")] #[test_case(CheckCode::SIM107, Path::new("SIM107.py"); "SIM107")] + #[test_case(CheckCode::SIM109, Path::new("SIM109.py"); "SIM109")] #[test_case(CheckCode::SIM110, Path::new("SIM110.py"); "SIM110")] #[test_case(CheckCode::SIM111, Path::new("SIM111.py"); "SIM111")] #[test_case(CheckCode::SIM117, Path::new("SIM117.py"); "SIM117")] diff --git a/src/flake8_simplify/plugins/ast_bool_op.rs b/src/flake8_simplify/plugins/ast_bool_op.rs index a583354101..ae3409dcb5 100644 --- a/src/flake8_simplify/plugins/ast_bool_op.rs +++ b/src/flake8_simplify/plugins/ast_bool_op.rs @@ -1,8 +1,9 @@ +use std::collections::BTreeMap; use std::iter; use itertools::Either::{Left, Right}; use rustc_hash::FxHashMap; -use rustpython_ast::{Boolop, Constant, Expr, ExprContext, ExprKind, Unaryop}; +use rustpython_ast::{Boolop, Cmpop, Constant, Expr, ExprContext, ExprKind, Unaryop}; use crate::ast::helpers::{create_expr, unparse_expr}; use crate::ast::types::Range; @@ -134,6 +135,72 @@ pub fn duplicate_isinstance_call(checker: &mut Checker, expr: &Expr) { } } +/// SIM109 +pub fn compare_with_tuple(checker: &mut Checker, expr: &Expr) { + let ExprKind::BoolOp { op: Boolop::Or, values } = &expr.node else { + return; + }; + + let mut id_to_values = BTreeMap::<&str, Vec<&Expr>>::new(); + for value in values { + let ExprKind::Compare { left, ops, comparators } = &value.node else { + continue; + }; + if ops.len() != 1 || comparators.len() != 1 { + continue; + } + if !matches!(&ops[0], Cmpop::Eq) { + continue; + } + let ExprKind::Name { id, .. } = &left.node else { + continue; + }; + let comparator = &comparators[0]; + if !matches!(&comparator.node, ExprKind::Name { .. }) { + continue; + } + id_to_values.entry(id).or_default().push(comparator); + } + + for (value, values) in id_to_values { + if values.len() == 1 { + continue; + } + let str_values = values + .iter() + .map(|value| unparse_expr(value, checker.style)) + .collect(); + let mut check = Check::new( + CheckKind::CompareWithTuple( + value.to_string(), + str_values, + unparse_expr(expr, checker.style), + ), + Range::from_located(expr), + ); + if checker.patch(&CheckCode::SIM109) { + // Create a `x in (a, b)` compare expr. + let in_expr = create_expr(ExprKind::Compare { + left: Box::new(create_expr(ExprKind::Name { + id: value.to_string(), + ctx: ExprContext::Load, + })), + ops: vec![Cmpop::In], + comparators: vec![create_expr(ExprKind::Tuple { + elts: values.into_iter().map(Clone::clone).collect(), + ctx: ExprContext::Load, + })], + }); + check.amend(Fix::replacement( + unparse_expr(&in_expr, checker.style), + expr.location, + expr.end_location.unwrap(), + )); + } + checker.add_check(check); + } +} + /// SIM220 pub fn a_and_not_a(checker: &mut Checker, expr: &Expr) { let ExprKind::BoolOp { op: Boolop::And, values, } = &expr.node else { diff --git a/src/flake8_simplify/plugins/mod.rs b/src/flake8_simplify/plugins/mod.rs index da812e3240..88b20f5496 100644 --- a/src/flake8_simplify/plugins/mod.rs +++ b/src/flake8_simplify/plugins/mod.rs @@ -1,4 +1,6 @@ -pub use ast_bool_op::{a_and_not_a, a_or_not_a, and_false, duplicate_isinstance_call, or_true}; +pub use ast_bool_op::{ + a_and_not_a, a_or_not_a, and_false, compare_with_tuple, duplicate_isinstance_call, or_true, +}; pub use ast_for::convert_loop_to_any_all; pub use ast_if::{nested_if_statements, use_ternary_operator}; pub use ast_with::multiple_with_statements; diff --git a/src/flake8_simplify/snapshots/ruff__flake8_simplify__tests__SIM109_SIM109.py.snap b/src/flake8_simplify/snapshots/ruff__flake8_simplify__tests__SIM109_SIM109.py.snap new file mode 100644 index 0000000000..e5e7cb1abe --- /dev/null +++ b/src/flake8_simplify/snapshots/ruff__flake8_simplify__tests__SIM109_SIM109.py.snap @@ -0,0 +1,26 @@ +--- +source: src/flake8_simplify/mod.rs +expression: checks +--- +- kind: + CompareWithTuple: + - a + - - b + - c + - a == b or a == c + location: + row: 2 + column: 3 + end_location: + row: 2 + column: 19 + fix: + content: "a in (b, c)" + location: + row: 2 + column: 3 + end_location: + row: 2 + column: 19 + parent: ~ + diff --git a/src/registry.rs b/src/registry.rs index 0d92c4c124..c81f12288d 100644 --- a/src/registry.rs +++ b/src/registry.rs @@ -223,6 +223,7 @@ pub enum CheckCode { SIM105, SIM107, SIM108, + SIM109, SIM110, SIM111, SIM117, @@ -963,6 +964,7 @@ pub enum CheckKind { SysVersionSlice1Referenced, // flake8-simplify UseTernaryOperator(String), + CompareWithTuple(String, Vec, String), DuplicateIsinstanceCall(String), AAndNotA(String), AOrNotA(String), @@ -1411,6 +1413,11 @@ impl CheckCode { CheckCode::SIM105 => CheckKind::UseContextlibSuppress("...".to_string()), CheckCode::SIM107 => CheckKind::ReturnInTryExceptFinally, CheckCode::SIM108 => CheckKind::UseTernaryOperator("..".to_string()), + CheckCode::SIM109 => CheckKind::CompareWithTuple( + "value".to_string(), + vec!["...".to_string(), "...".to_string()], + "value == ... or value == ...".to_string(), + ), CheckCode::SIM110 => { CheckKind::ConvertLoopToAny("return any(x for x in y)".to_string()) } @@ -1961,6 +1968,7 @@ impl CheckCode { CheckCode::SIM105 => CheckCategory::Flake8Simplify, CheckCode::SIM107 => CheckCategory::Flake8Simplify, CheckCode::SIM108 => CheckCategory::Flake8Simplify, + CheckCode::SIM109 => CheckCategory::Flake8Simplify, CheckCode::SIM110 => CheckCategory::Flake8Simplify, CheckCode::SIM111 => CheckCategory::Flake8Simplify, CheckCode::SIM117 => CheckCategory::Flake8Simplify, @@ -2216,12 +2224,13 @@ impl CheckKind { CheckKind::SysVersionCmpStr10 => &CheckCode::YTT302, CheckKind::SysVersionSlice1Referenced => &CheckCode::YTT303, // flake8-simplify - CheckKind::DuplicateIsinstanceCall(..) => &CheckCode::SIM101, CheckKind::AAndNotA(..) => &CheckCode::SIM220, CheckKind::AOrNotA(..) => &CheckCode::SIM221, CheckKind::AndFalse => &CheckCode::SIM223, + CheckKind::CompareWithTuple(..) => &CheckCode::SIM109, CheckKind::ConvertLoopToAll(..) => &CheckCode::SIM111, CheckKind::ConvertLoopToAny(..) => &CheckCode::SIM110, + CheckKind::DuplicateIsinstanceCall(..) => &CheckCode::SIM101, CheckKind::KeyInDict(..) => &CheckCode::SIM118, CheckKind::MultipleWithStatements => &CheckCode::SIM117, CheckKind::NestedIfStatements => &CheckCode::SIM102, @@ -2987,6 +2996,10 @@ impl CheckKind { "`sys.version[:1]` referenced (python10), use `sys.version_info`".to_string() } // flake8-simplify + CheckKind::CompareWithTuple(value, values, or_op) => { + let values = values.join(", "); + format!("Use `{value} in ({values})` instead of `{or_op}`") + } CheckKind::DuplicateIsinstanceCall(name) => { format!("Multiple `isinstance` calls for `{name}`, merge into a single call") } @@ -3663,6 +3676,7 @@ impl CheckKind { | CheckKind::BlankLineBeforeSection(..) | CheckKind::CapitalizeSectionName(..) | CheckKind::CommentedOutCode + | CheckKind::CompareWithTuple(..) | CheckKind::ConvertLoopToAll(..) | CheckKind::ConvertLoopToAny(..) | CheckKind::ConvertNamedTupleFunctionalToClass(..) @@ -3798,6 +3812,10 @@ impl CheckKind { } CheckKind::CapitalizeSectionName(name) => Some(format!("Capitalize \"{name}\"")), CheckKind::CommentedOutCode => Some("Remove commented-out code".to_string()), + CheckKind::CompareWithTuple(value, values, or_op) => { + let values = values.join(", "); + Some(format!("Replace `{or_op}` with `{value} in {values}`")) + } CheckKind::ConvertLoopToAll(all) => Some(format!("Replace with `{all}`")), CheckKind::ConvertLoopToAny(any) => Some(format!("Replace with `{any}`")), CheckKind::ConvertTypedDictFunctionalToClass(name)