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

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;