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 00000000000000..11695c4273fa1e --- /dev/null +++ b/crates/ruff_linter/resources/test/fixtures/refurb/FURB142.py @@ -0,0 +1,56 @@ +# 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) + + +# 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 a5b7cc2cd8972f..0b7839767356d8 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 efbf51abf39d04..7b5dbd8a24502e 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 f67ddee937a017..1408109b2c6e63 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 00000000000000..4b9646c774fbed --- /dev/null +++ b/crates/ruff_linter/src/rules/refurb/rules/for_loop_set_mutations.rs @@ -0,0 +1,125 @@ +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 updating list with iterable object by calling `add`/`discard` on each element +/// separately in the `for` loop. +/// +/// ## Why is this bad? +/// When adding/removing a batch of elements to/from a set, it's more readable to call +/// an appropriate method, instead of add/remove 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 a395e6065a71a5..832342c509e13b 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 00000000000000..e385a739b28358 --- /dev/null +++ b/crates/ruff_linter/src/rules/refurb/snapshots/ruff_linter__rules__refurb__tests__FURB142_FURB142.py.snap @@ -0,0 +1,161 @@ +--- +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 + | + = 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 | +29 28 | # False negative diff --git a/ruff.schema.json b/ruff.schema.json index ad659af9f05fd7..a0982b2c876908 100644 --- a/ruff.schema.json +++ b/ruff.schema.json @@ -3014,6 +3014,7 @@ "FURB136", "FURB14", "FURB140", + "FURB142", "FURB145", "FURB148", "FURB15",