diff --git a/crates/ruff_linter/resources/test/fixtures/refurb/FURB142.py b/crates/ruff_linter/resources/test/fixtures/refurb/FURB142.py new file mode 100644 index 0000000000..6a4fd0521b --- /dev/null +++ b/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) diff --git a/crates/ruff_linter/src/checkers/ast/analyze/statement.rs b/crates/ruff_linter/src/checkers/ast/analyze/statement.rs index a5b7cc2cd8..0b78397673 100644 --- a/crates/ruff_linter/src/checkers/ast/analyze/statement.rs +++ b/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 { diff --git a/crates/ruff_linter/src/codes.rs b/crates/ruff_linter/src/codes.rs index efbf51abf3..7b5dbd8a24 100644 --- a/crates/ruff_linter/src/codes.rs +++ b/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), diff --git a/crates/ruff_linter/src/rules/refurb/mod.rs b/crates/ruff_linter/src/rules/refurb/mod.rs index f67ddee937..1408109b2c 100644 --- a/crates/ruff_linter/src/rules/refurb/mod.rs +++ b/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"))] diff --git a/crates/ruff_linter/src/rules/refurb/rules/for_loop_set_mutations.rs b/crates/ruff_linter/src/rules/refurb/rules/for_loop_set_mutations.rs new file mode 100644 index 0000000000..b019c6bd3c --- /dev/null +++ b/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, + ))), + ); +} diff --git a/crates/ruff_linter/src/rules/refurb/rules/mod.rs b/crates/ruff_linter/src/rules/refurb/rules/mod.rs index a395e6065a..832342c509 100644 --- a/crates/ruff_linter/src/rules/refurb/rules/mod.rs +++ b/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; diff --git a/crates/ruff_linter/src/rules/refurb/snapshots/ruff_linter__rules__refurb__tests__FURB142_FURB142.py.snap b/crates/ruff_linter/src/rules/refurb/snapshots/ruff_linter__rules__refurb__tests__FURB142_FURB142.py.snap new file mode 100644 index 0000000000..cd781a8bbe --- /dev/null +++ b/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 | diff --git a/ruff.schema.json b/ruff.schema.json index ad659af9f0..a0982b2c87 100644 --- a/ruff.schema.json +++ b/ruff.schema.json @@ -3014,6 +3014,7 @@ "FURB136", "FURB14", "FURB140", + "FURB142", "FURB145", "FURB148", "FURB15",