Preserve unmatched comparators in `SIM109` (#1998)

Closes #1993.

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 `...` | 🛠 |

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
+    d
+
+# OK
+if a == b or a == c():
+    d
+
+# OK
+if (
+    a == b
+    # This comment prevents us from raising SIM109
+    or a == c
+):
+    d

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<Expr> = 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,

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: ~
 

src/violations.rs

@@ -2894,27 +2894,25 @@ impl AlwaysAutofixableViolation for UseTernaryOperator {
 }
 
 define_violation!(
-    pub struct CompareWithTuple(pub String, pub Vec<String>, 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(),
+        }
     }
 }