diff --git a/README.md b/README.md index 54bac0199b..ea51ab07cc 100644 --- a/README.md +++ b/README.md @@ -1009,7 +1009,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 == ...` | 🛠 | +| SIM109 | CompareWithTuple | Use `value in (... ,...)` instead of multiple equality comparisons | 🛠 | | 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 | 🛠 | | SIM112 | UseCapitalEnvironmentVariables | Use capitalized environment variable `...` instead of `...` | 🛠 | diff --git a/resources/test/fixtures/flake8_simplify/SIM109.py b/resources/test/fixtures/flake8_simplify/SIM109.py index c794e6fd84..f28c293a72 100644 --- a/resources/test/fixtures/flake8_simplify/SIM109.py +++ b/resources/test/fixtures/flake8_simplify/SIM109.py @@ -1,7 +1,31 @@ -# Bad +# SIM109 if a == b or a == c: d -# Good +# SIM109 +if (a == b or a == c) and None: + d + +# SIM109 +if a == b or a == c or None: + d + +# SIM109 +if a == b or None or a == c: + d + +# OK if a in (b, c): - d \ No newline at end of file + d + +# OK +if a == b or a == c(): + d + +# OK +if ( + a == b + # This comment prevents us from raising SIM109 + or a == c +): + d diff --git a/src/rules/flake8_simplify/rules/ast_bool_op.rs b/src/rules/flake8_simplify/rules/ast_bool_op.rs index 5002247d2d..b576d59f0c 100644 --- a/src/rules/flake8_simplify/rules/ast_bool_op.rs +++ b/src/rules/flake8_simplify/rules/ast_bool_op.rs @@ -2,10 +2,9 @@ use std::collections::BTreeMap; use std::iter; use itertools::Either::{Left, Right}; -use rustc_hash::FxHashMap; use rustpython_ast::{Boolop, Cmpop, Constant, Expr, ExprContext, ExprKind, Unaryop}; -use crate::ast::helpers::{contains_effect, create_expr, unparse_expr}; +use crate::ast::helpers::{contains_effect, create_expr, has_comments_in, unparse_expr}; use crate::ast::types::Range; use crate::checkers::ast::Checker; use crate::fix::Fix; @@ -30,7 +29,7 @@ pub fn duplicate_isinstance_call(checker: &mut Checker, expr: &Expr) { // Locate duplicate `isinstance` calls, represented as a map from argument name // to indices of the relevant `Expr` instances in `values`. - let mut duplicates = FxHashMap::default(); + let mut duplicates = BTreeMap::new(); for (index, call) in values.iter().enumerate() { // Verify that this is an `isinstance` call. let ExprKind::Call { func, args, keywords } = &call.node else { @@ -136,62 +135,98 @@ pub fn duplicate_isinstance_call(checker: &mut Checker, expr: &Expr) { } } +fn match_eq_target(expr: &Expr) -> Option<(&str, &Expr)> { + let ExprKind::Compare { left, ops, comparators } = &expr.node else { + return None; + }; + if ops.len() != 1 || comparators.len() != 1 { + return None; + } + if !matches!(&ops[0], Cmpop::Eq) { + return None; + } + let ExprKind::Name { id, .. } = &left.node else { + return None; + }; + let comparator = &comparators[0]; + if !matches!(&comparator.node, ExprKind::Name { .. }) { + return None; + } + Some((id, comparator)) +} + /// 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; + // Given `a == "foo" or a == "bar"`, we generate `{"a": [(0, "foo"), (1, + // "bar")]}`. + let mut id_to_comparators: BTreeMap<&str, Vec<(usize, &Expr)>> = BTreeMap::new(); + for (index, value) in values.iter().enumerate() { + if let Some((id, comparator)) = match_eq_target(value) { + id_to_comparators + .entry(id) + .or_default() + .push((index, comparator)); } - 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 { + for (id, matches) in id_to_comparators { + if matches.len() == 1 { continue; } - let str_values = values + + let (indices, comparators): (Vec<_>, Vec<_>) = matches.iter().copied().unzip(); + + // Avoid rewriting (e.g.) `a == "foo" or a == f()`. + if comparators .iter() - .map(|value| unparse_expr(value, checker.stylist)) - .collect(); + .any(|expr| contains_effect(checker, expr)) + { + continue; + } + + // Avoid removing comments. + if has_comments_in(Range::from_located(expr), checker.locator) { + continue; + } + + // Create a `x in (a, b)` expression. + let in_expr = create_expr(ExprKind::Compare { + left: Box::new(create_expr(ExprKind::Name { + id: id.to_string(), + ctx: ExprContext::Load, + })), + ops: vec![Cmpop::In], + comparators: vec![create_expr(ExprKind::Tuple { + elts: comparators.into_iter().map(Clone::clone).collect(), + ctx: ExprContext::Load, + })], + }); let mut diagnostic = Diagnostic::new( - violations::CompareWithTuple( - value.to_string(), - str_values, - unparse_expr(expr, checker.stylist), - ), + violations::CompareWithTuple { + replacement: unparse_expr(&in_expr, checker.stylist), + }, Range::from_located(expr), ); if checker.patch(&Rule::CompareWithTuple) { - // 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, - })], - }); + let unmatched: Vec = values + .iter() + .enumerate() + .filter(|(index, _)| !indices.contains(index)) + .map(|(_, elt)| elt.clone()) + .collect(); + let in_expr = if unmatched.is_empty() { + in_expr + } else { + // Wrap in a `x in (a, b) or ...` boolean operation. + create_expr(ExprKind::BoolOp { + op: Boolop::Or, + values: iter::once(in_expr).chain(unmatched).collect(), + }) + }; diagnostic.amend(Fix::replacement( unparse_expr(&in_expr, checker.stylist), expr.location, diff --git a/src/rules/flake8_simplify/snapshots/ruff__rules__flake8_simplify__tests__SIM109_SIM109.py.snap b/src/rules/flake8_simplify/snapshots/ruff__rules__flake8_simplify__tests__SIM109_SIM109.py.snap index ffc6d77caf..a7fa8f985d 100644 --- a/src/rules/flake8_simplify/snapshots/ruff__rules__flake8_simplify__tests__SIM109_SIM109.py.snap +++ b/src/rules/flake8_simplify/snapshots/ruff__rules__flake8_simplify__tests__SIM109_SIM109.py.snap @@ -4,10 +4,7 @@ expression: diagnostics --- - kind: CompareWithTuple: - - a - - - b - - c - - a == b or a == c + replacement: "a in (b, c)" location: row: 2 column: 3 @@ -23,4 +20,58 @@ expression: diagnostics row: 2 column: 19 parent: ~ +- kind: + CompareWithTuple: + replacement: "a in (b, c)" + location: + row: 6 + column: 4 + end_location: + row: 6 + column: 20 + fix: + content: "a in (b, c)" + location: + row: 6 + column: 4 + end_location: + row: 6 + column: 20 + parent: ~ +- kind: + CompareWithTuple: + replacement: "a in (b, c)" + location: + row: 10 + column: 3 + end_location: + row: 10 + column: 27 + fix: + content: "a in (b, c) or None" + location: + row: 10 + column: 3 + end_location: + row: 10 + column: 27 + parent: ~ +- kind: + CompareWithTuple: + replacement: "a in (b, c)" + location: + row: 14 + column: 3 + end_location: + row: 14 + column: 27 + fix: + content: "a in (b, c) or None" + location: + row: 14 + column: 3 + end_location: + row: 14 + column: 27 + parent: ~ diff --git a/src/violations.rs b/src/violations.rs index 292e6af241..061ff7af29 100644 --- a/src/violations.rs +++ b/src/violations.rs @@ -2894,27 +2894,25 @@ impl AlwaysAutofixableViolation for UseTernaryOperator { } define_violation!( - pub struct CompareWithTuple(pub String, pub Vec, pub String); + pub struct CompareWithTuple { + pub replacement: String, + } ); impl AlwaysAutofixableViolation for CompareWithTuple { fn message(&self) -> String { - let CompareWithTuple(value, values, or_op) = self; - let values = values.join(", "); - format!("Use `{value} in ({values})` instead of `{or_op}`") + let CompareWithTuple { replacement } = self; + format!("Use `{replacement}` instead of multiple equality comparisons") } fn autofix_title(&self) -> String { - let CompareWithTuple(value, values, or_op) = self; - let values = values.join(", "); - format!("Replace `{or_op}` with `{value} in {values}`") + let CompareWithTuple { replacement, .. } = self; + format!("Replace with `{replacement}`") } fn placeholder() -> Self { - CompareWithTuple( - "value".to_string(), - vec!["...".to_string(), "...".to_string()], - "value == ... or value == ...".to_string(), - ) + CompareWithTuple { + replacement: "value in (... ,...)".to_string(), + } } }