[refurb] Implement for-loop-set-mutations (FURB142) (#10583)

Co-authored-by: Alex Waygood <alex.waygood@gmail.com>

crates/ruff_linter/resources/test/fixtures/refurb/FURB142.py

@@ -0,0 +1,61 @@
+# Errors
+
+s = set()
+
+for x in [1, 2, 3]:
+    s.add(x)
+
+for x in {1, 2, 3}:
+    s.add(x)
+
+for x in (1, 2, 3):
+    s.add(x)
+
+for x in (1, 2, 3):
+    s.discard(x)
+
+for x in (1, 2, 3):
+    s.add(x + 1)
+
+for x, y in ((1, 2), (3, 4)):
+    s.add((x, y))
+
+num = 123
+
+for x in (1, 2, 3):
+    s.add(num)
+
+for x in (1, 2, 3):
+    s.add((num, x))
+
+for x in (1, 2, 3):
+    s.add(x + num)
+
+# False negative
+
+class C:
+    s: set[int]
+
+
+c = C()
+for x in (1, 2, 3):
+    c.s.add(x)
+
+# Ok
+
+s.update(x for x in (1, 2, 3))
+
+for x in (1, 2, 3):
+    s.add(x)
+else:
+    pass
+
+
+async def f(y):
+    async for x in y:
+        s.add(x)
+
+
+def g():
+    for x in (set(),):
+        x.add(x)

crates/ruff_linter/src/checkers/ast/analyze/statement.rs

@@ -1315,6 +1315,9 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
                 if checker.enabled(Rule::TryExceptInLoop) {
                     perflint::rules::try_except_in_loop(checker, body);
                 }
+                if checker.enabled(Rule::ForLoopSetMutations) {
+                    refurb::rules::for_loop_set_mutations(checker, for_stmt);
+                }
             }
         }
         Stmt::Try(ast::StmtTry {

crates/ruff_linter/src/codes.rs

@@ -1044,6 +1044,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> {
         (Refurb, "132") => (RuleGroup::Nursery, rules::refurb::rules::CheckAndRemoveFromSet),
         (Refurb, "136") => (RuleGroup::Preview, rules::refurb::rules::IfExprMinMax),
         (Refurb, "140") => (RuleGroup::Preview, rules::refurb::rules::ReimplementedStarmap),
+        (Refurb, "142") => (RuleGroup::Preview, rules::refurb::rules::ForLoopSetMutations),
         (Refurb, "145") => (RuleGroup::Preview, rules::refurb::rules::SliceCopy),
         (Refurb, "148") => (RuleGroup::Preview, rules::refurb::rules::UnnecessaryEnumerate),
         (Refurb, "152") => (RuleGroup::Preview, rules::refurb::rules::MathConstant),

crates/ruff_linter/src/rules/refurb/mod.rs

@@ -22,6 +22,7 @@ mod tests {
     #[test_case(Rule::CheckAndRemoveFromSet, Path::new("FURB132.py"))]
     #[test_case(Rule::IfExprMinMax, Path::new("FURB136.py"))]
     #[test_case(Rule::ReimplementedStarmap, Path::new("FURB140.py"))]
+    #[test_case(Rule::ForLoopSetMutations, Path::new("FURB142.py"))]
     #[test_case(Rule::SliceCopy, Path::new("FURB145.py"))]
     #[test_case(Rule::UnnecessaryEnumerate, Path::new("FURB148.py"))]
     #[test_case(Rule::MathConstant, Path::new("FURB152.py"))]

crates/ruff_linter/src/rules/refurb/rules/for_loop_set_mutations.rs

@@ -0,0 +1,127 @@
+use ruff_diagnostics::{AlwaysFixableViolation, Diagnostic, Edit, Fix};
+use ruff_macros::{derive_message_formats, violation};
+use ruff_python_ast::{Expr, Stmt, StmtFor};
+use ruff_python_semantic::analyze::typing;
+
+use crate::checkers::ast::Checker;
+
+/// ## What it does
+/// Checks for code that updates a set with the contents of an iterable by
+/// using a `for` loop to call `.add()` or `.discard()` on each element
+/// separately.
+///
+/// ## Why is this bad?
+/// When adding or removing a batch of elements to or from a set, it's more
+/// idiomatic to use a single method call rather than adding or removing
+/// elements one by one.
+///
+/// ## Example
+/// ```python
+/// s = set()
+///
+/// for x in (1, 2, 3):
+///     s.add(x)
+///
+/// for x in (1, 2, 3):
+///     s.discard(x)
+/// ```
+///
+/// Use instead:
+/// ```python
+/// s = set()
+///
+/// s.update((1, 2, 3))
+/// s.difference_update((1, 2, 3))
+/// ```
+///
+/// ## References
+/// - [Python documentation: `set`](https://docs.python.org/3/library/stdtypes.html#set)
+#[violation]
+pub struct ForLoopSetMutations {
+    method_name: &'static str,
+    batch_method_name: &'static str,
+}
+
+impl AlwaysFixableViolation for ForLoopSetMutations {
+    #[derive_message_formats]
+    fn message(&self) -> String {
+        format!("Use of `set.{}()` in a for loop", self.method_name)
+    }
+
+    fn fix_title(&self) -> String {
+        format!("Replace with `.{}()`", self.batch_method_name)
+    }
+}
+
+// FURB142
+pub(crate) fn for_loop_set_mutations(checker: &mut Checker, for_stmt: &StmtFor) {
+    if !for_stmt.orelse.is_empty() {
+        return;
+    }
+    let [Stmt::Expr(stmt_expr)] = for_stmt.body.as_slice() else {
+        return;
+    };
+    let Expr::Call(expr_call) = stmt_expr.value.as_ref() else {
+        return;
+    };
+    let Expr::Attribute(expr_attr) = expr_call.func.as_ref() else {
+        return;
+    };
+    if !expr_call.arguments.keywords.is_empty() {
+        return;
+    }
+
+    let (method_name, batch_method_name) = match expr_attr.attr.as_str() {
+        "add" => ("add", "update"),
+        "discard" => ("discard", "difference_update"),
+        _ => {
+            return;
+        }
+    };
+
+    let Expr::Name(set) = expr_attr.value.as_ref() else {
+        return;
+    };
+
+    if !checker
+        .semantic()
+        .resolve_name(set)
+        .is_some_and(|s| typing::is_set(checker.semantic().binding(s), checker.semantic()))
+    {
+        return;
+    }
+    let [arg] = expr_call.arguments.args.as_ref() else {
+        return;
+    };
+
+    let content = match (for_stmt.target.as_ref(), arg) {
+        (Expr::Name(for_target), Expr::Name(arg)) if for_target.id == arg.id => {
+            format!(
+                "{}.{batch_method_name}({})",
+                set.id,
+                checker.locator().slice(for_stmt.iter.as_ref())
+            )
+        }
+        (for_target, arg) => format!(
+            "{}.{batch_method_name}({} for {} in {})",
+            set.id,
+            checker.locator().slice(arg),
+            checker.locator().slice(for_target),
+            checker.locator().slice(for_stmt.iter.as_ref())
+        ),
+    };
+
+    checker.diagnostics.push(
+        Diagnostic::new(
+            ForLoopSetMutations {
+                method_name,
+                batch_method_name,
+            },
+            for_stmt.range,
+        )
+        .with_fix(Fix::safe_edit(Edit::range_replacement(
+            content,
+            for_stmt.range,
+        ))),
+    );
+}

crates/ruff_linter/src/rules/refurb/rules/mod.rs

@@ -1,6 +1,7 @@
 pub(crate) use bit_count::*;
 pub(crate) use check_and_remove_from_set::*;
 pub(crate) use delete_full_slice::*;
+pub(crate) use for_loop_set_mutations::*;
 pub(crate) use hashlib_digest_hex::*;
 pub(crate) use if_expr_min_max::*;
 pub(crate) use implicit_cwd::*;
@@ -25,6 +26,7 @@ pub(crate) use verbose_decimal_constructor::*;
 mod bit_count;
 mod check_and_remove_from_set;
 mod delete_full_slice;
+mod for_loop_set_mutations;
 mod hashlib_digest_hex;
 mod if_expr_min_max;
 mod implicit_cwd;

crates/ruff_linter/src/rules/refurb/snapshots/ruff_linter__rules__refurb__tests__FURB142_FURB142.py.snap

@@ -0,0 +1,209 @@
+---
+source: crates/ruff_linter/src/rules/refurb/mod.rs
+---
+FURB142.py:5:1: FURB142 [*] Use of `set.add()` in a for loop
+  |
+3 |   s = set()
+4 |   
+5 | / for x in [1, 2, 3]:
+6 | |     s.add(x)
+  | |____________^ FURB142
+7 |   
+8 |   for x in {1, 2, 3}:
+  |
+  = help: Replace with `.update()`
+
+ℹ Safe fix
+2 2 | 
+3 3 | s = set()
+4 4 | 
+5   |-for x in [1, 2, 3]:
+6   |-    s.add(x)
+  5 |+s.update([1, 2, 3])
+7 6 | 
+8 7 | for x in {1, 2, 3}:
+9 8 |     s.add(x)
+
+FURB142.py:8:1: FURB142 [*] Use of `set.add()` in a for loop
+   |
+ 6 |       s.add(x)
+ 7 |   
+ 8 | / for x in {1, 2, 3}:
+ 9 | |     s.add(x)
+   | |____________^ FURB142
+10 |   
+11 |   for x in (1, 2, 3):
+   |
+   = help: Replace with `.update()`
+
+ℹ Safe fix
+5  5  | for x in [1, 2, 3]:
+6  6  |     s.add(x)
+7  7  | 
+8     |-for x in {1, 2, 3}:
+9     |-    s.add(x)
+   8  |+s.update({1, 2, 3})
+10 9  | 
+11 10 | for x in (1, 2, 3):
+12 11 |     s.add(x)
+
+FURB142.py:11:1: FURB142 [*] Use of `set.add()` in a for loop
+   |
+ 9 |       s.add(x)
+10 |   
+11 | / for x in (1, 2, 3):
+12 | |     s.add(x)
+   | |____________^ FURB142
+13 |   
+14 |   for x in (1, 2, 3):
+   |
+   = help: Replace with `.update()`
+
+ℹ Safe fix
+8  8  | for x in {1, 2, 3}:
+9  9  |     s.add(x)
+10 10 | 
+11    |-for x in (1, 2, 3):
+12    |-    s.add(x)
+   11 |+s.update((1, 2, 3))
+13 12 | 
+14 13 | for x in (1, 2, 3):
+15 14 |     s.discard(x)
+
+FURB142.py:14:1: FURB142 [*] Use of `set.discard()` in a for loop
+   |
+12 |       s.add(x)
+13 |   
+14 | / for x in (1, 2, 3):
+15 | |     s.discard(x)
+   | |________________^ FURB142
+16 |   
+17 |   for x in (1, 2, 3):
+   |
+   = help: Replace with `.difference_update()`
+
+ℹ Safe fix
+11 11 | for x in (1, 2, 3):
+12 12 |     s.add(x)
+13 13 | 
+14    |-for x in (1, 2, 3):
+15    |-    s.discard(x)
+   14 |+s.difference_update((1, 2, 3))
+16 15 | 
+17 16 | for x in (1, 2, 3):
+18 17 |     s.add(x + 1)
+
+FURB142.py:17:1: FURB142 [*] Use of `set.add()` in a for loop
+   |
+15 |       s.discard(x)
+16 |   
+17 | / for x in (1, 2, 3):
+18 | |     s.add(x + 1)
+   | |________________^ FURB142
+19 |   
+20 |   for x, y in ((1, 2), (3, 4)):
+   |
+   = help: Replace with `.update()`
+
+ℹ Safe fix
+14 14 | for x in (1, 2, 3):
+15 15 |     s.discard(x)
+16 16 | 
+17    |-for x in (1, 2, 3):
+18    |-    s.add(x + 1)
+   17 |+s.update(x + 1 for x in (1, 2, 3))
+19 18 | 
+20 19 | for x, y in ((1, 2), (3, 4)):
+21 20 |     s.add((x, y))
+
+FURB142.py:20:1: FURB142 [*] Use of `set.add()` in a for loop
+   |
+18 |       s.add(x + 1)
+19 |   
+20 | / for x, y in ((1, 2), (3, 4)):
+21 | |     s.add((x, y))
+   | |_________________^ FURB142
+22 |   
+23 |   num = 123
+   |
+   = help: Replace with `.update()`
+
+ℹ Safe fix
+17 17 | for x in (1, 2, 3):
+18 18 |     s.add(x + 1)
+19 19 | 
+20    |-for x, y in ((1, 2), (3, 4)):
+21    |-    s.add((x, y))
+   20 |+s.update((x, y) for x, y in ((1, 2), (3, 4)))
+22 21 | 
+23 22 | num = 123
+24 23 | 
+
+FURB142.py:25:1: FURB142 [*] Use of `set.add()` in a for loop
+   |
+23 |   num = 123
+24 |   
+25 | / for x in (1, 2, 3):
+26 | |     s.add(num)
+   | |______________^ FURB142
+27 |   
+28 |   for x in (1, 2, 3):
+   |
+   = help: Replace with `.update()`
+
+ℹ Safe fix
+22 22 | 
+23 23 | num = 123
+24 24 | 
+25    |-for x in (1, 2, 3):
+26    |-    s.add(num)
+   25 |+s.update(num for x in (1, 2, 3))
+27 26 | 
+28 27 | for x in (1, 2, 3):
+29 28 |     s.add((num, x))
+
+FURB142.py:28:1: FURB142 [*] Use of `set.add()` in a for loop
+   |
+26 |       s.add(num)
+27 |   
+28 | / for x in (1, 2, 3):
+29 | |     s.add((num, x))
+   | |___________________^ FURB142
+30 |   
+31 |   for x in (1, 2, 3):
+   |
+   = help: Replace with `.update()`
+
+ℹ Safe fix
+25 25 | for x in (1, 2, 3):
+26 26 |     s.add(num)
+27 27 | 
+28    |-for x in (1, 2, 3):
+29    |-    s.add((num, x))
+   28 |+s.update((num, x) for x in (1, 2, 3))
+30 29 | 
+31 30 | for x in (1, 2, 3):
+32 31 |     s.add(x + num)
+
+FURB142.py:31:1: FURB142 [*] Use of `set.add()` in a for loop
+   |
+29 |       s.add((num, x))
+30 |   
+31 | / for x in (1, 2, 3):
+32 | |     s.add(x + num)
+   | |__________________^ FURB142
+33 |   
+34 |   # False negative
+   |
+   = help: Replace with `.update()`
+
+ℹ Safe fix
+28 28 | for x in (1, 2, 3):
+29 29 |     s.add((num, x))
+30 30 | 
+31    |-for x in (1, 2, 3):
+32    |-    s.add(x + num)
+   31 |+s.update(x + num for x in (1, 2, 3))
+33 32 | 
+34 33 | # False negative
+35 34 |

ruff.schema.json

@@ -3014,6 +3014,7 @@
         "FURB136",
         "FURB14",
         "FURB140",
+        "FURB142",
         "FURB145",
         "FURB148",
         "FURB15",