Skip to content

Commit 05a2132

Browse files
alex-700alex-700
committed
[refurb] Implement manual_set_update (FURB142)
1 parent 877a914 commit 05a2132

File tree

8 files changed

+350
-0
lines changed

8 files changed

+350
-0
lines changed

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

+56
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
# Errors
2+
3+
s = set()
4+
5+
for x in [1, 2, 3]:
6+
s.add(x)
7+
8+
for x in {1, 2, 3}:
9+
s.add(x)
10+
11+
for x in (1, 2, 3):
12+
s.add(x)
13+
14+
for x in (1, 2, 3):
15+
s.discard(x)
16+
17+
for x in (1, 2, 3):
18+
s.add(x + 1)
19+
20+
for x, y in ((1, 2), (3, 4)):
21+
s.add((x, y))
22+
23+
num = 123
24+
25+
for x in (1, 2, 3):
26+
s.add(num)
27+
28+
29+
# False negative
30+
31+
class C:
32+
s: set[int]
33+
34+
35+
c = C()
36+
for x in (1, 2, 3):
37+
c.s.add(x)
38+
39+
# Ok
40+
41+
s.update(x for x in (1, 2, 3))
42+
43+
for x in (1, 2, 3):
44+
s.add(x)
45+
else:
46+
pass
47+
48+
49+
async def f(y):
50+
async for x in y:
51+
s.add(x)
52+
53+
54+
def g():
55+
for x in (set(),):
56+
x.add(x)

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

+3
Original file line numberDiff line numberDiff line change
@@ -1315,6 +1315,9 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
13151315
if checker.enabled(Rule::TryExceptInLoop) {
13161316
perflint::rules::try_except_in_loop(checker, body);
13171317
}
1318+
if checker.enabled(Rule::ForLoopSetMutations) {
1319+
refurb::rules::for_loop_set_mutations(checker, for_stmt);
1320+
}
13181321
}
13191322
}
13201323
Stmt::Try(ast::StmtTry {

‎crates/ruff_linter/src/codes.rs

+1
Original file line numberDiff line numberDiff line change
@@ -1044,6 +1044,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> {
10441044
(Refurb, "132") => (RuleGroup::Nursery, rules::refurb::rules::CheckAndRemoveFromSet),
10451045
(Refurb, "136") => (RuleGroup::Preview, rules::refurb::rules::IfExprMinMax),
10461046
(Refurb, "140") => (RuleGroup::Preview, rules::refurb::rules::ReimplementedStarmap),
1047+
(Refurb, "142") => (RuleGroup::Preview, rules::refurb::rules::ForLoopSetMutations),
10471048
(Refurb, "145") => (RuleGroup::Preview, rules::refurb::rules::SliceCopy),
10481049
(Refurb, "148") => (RuleGroup::Preview, rules::refurb::rules::UnnecessaryEnumerate),
10491050
(Refurb, "152") => (RuleGroup::Preview, rules::refurb::rules::MathConstant),

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

+1
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ mod tests {
2222
#[test_case(Rule::CheckAndRemoveFromSet, Path::new("FURB132.py"))]
2323
#[test_case(Rule::IfExprMinMax, Path::new("FURB136.py"))]
2424
#[test_case(Rule::ReimplementedStarmap, Path::new("FURB140.py"))]
25+
#[test_case(Rule::ForLoopSetMutations, Path::new("FURB142.py"))]
2526
#[test_case(Rule::SliceCopy, Path::new("FURB145.py"))]
2627
#[test_case(Rule::UnnecessaryEnumerate, Path::new("FURB148.py"))]
2728
#[test_case(Rule::MathConstant, Path::new("FURB152.py"))]

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

+125
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,125 @@
1+
use ruff_diagnostics::{AlwaysFixableViolation, Diagnostic, Edit, Fix};
2+
use ruff_macros::{derive_message_formats, violation};
3+
use ruff_python_ast::{Expr, Stmt, StmtFor};
4+
use ruff_python_semantic::analyze::typing;
5+
6+
use crate::checkers::ast::Checker;
7+
8+
/// ## What it does
9+
/// Checks for updating list with iterable object by calling `add`/`discard` on each element
10+
/// separately in the `for` loop.
11+
///
12+
/// ## Why is this bad?
13+
/// When adding/removing a batch of elements to/from a set, it's more readable to call
14+
/// an appropriate method, instead of add/remove elements one-by-one.
15+
///
16+
/// ## Example
17+
/// ```python
18+
/// s = set()
19+
///
20+
/// for x in (1, 2, 3):
21+
/// s.add(x)
22+
///
23+
/// for x in (1, 2, 3):
24+
/// s.discard(x)
25+
/// ```
26+
///
27+
/// Use instead:
28+
/// ```python
29+
/// s = set()
30+
///
31+
/// s.update((1, 2, 3))
32+
/// s.difference_update((1, 2, 3))
33+
/// ```
34+
///
35+
/// ## References
36+
/// - [Python documentation: `set`](https://docs.python.org/3/library/stdtypes.html#set)
37+
#[violation]
38+
pub struct ForLoopSetMutations {
39+
method_name: &'static str,
40+
batch_method_name: &'static str,
41+
}
42+
43+
impl AlwaysFixableViolation for ForLoopSetMutations {
44+
#[derive_message_formats]
45+
fn message(&self) -> String {
46+
format!("Use of set.{} in a for loop", self.method_name)
47+
}
48+
49+
fn fix_title(&self) -> String {
50+
format!("Replace with `{}`", self.batch_method_name)
51+
}
52+
}
53+
54+
// FURB142
55+
pub(crate) fn for_loop_set_mutations(checker: &mut Checker, for_stmt: &StmtFor) {
56+
if !for_stmt.orelse.is_empty() {
57+
return;
58+
}
59+
let [Stmt::Expr(stmt_expr)] = for_stmt.body.as_slice() else {
60+
return;
61+
};
62+
let Expr::Call(expr_call) = stmt_expr.value.as_ref() else {
63+
return;
64+
};
65+
let Expr::Attribute(expr_attr) = expr_call.func.as_ref() else {
66+
return;
67+
};
68+
if !expr_call.arguments.keywords.is_empty() {
69+
return;
70+
}
71+
72+
let (method_name, batch_method_name) = match expr_attr.attr.as_str() {
73+
"add" => ("add", "update"),
74+
"discard" => ("discard", "difference_update"),
75+
_ => {
76+
return;
77+
}
78+
};
79+
80+
let Expr::Name(set) = expr_attr.value.as_ref() else {
81+
return;
82+
};
83+
84+
if !checker
85+
.semantic()
86+
.resolve_name(set)
87+
.is_some_and(|s| typing::is_set(checker.semantic().binding(s), checker.semantic()))
88+
{
89+
return;
90+
}
91+
let [arg] = expr_call.arguments.args.as_ref() else {
92+
return;
93+
};
94+
95+
let content = match (for_stmt.target.as_ref(), arg) {
96+
(Expr::Name(for_target), Expr::Name(arg)) if for_target.id == arg.id => {
97+
format!(
98+
"{}.{batch_method_name}({})",
99+
set.id,
100+
checker.locator().slice(for_stmt.iter.as_ref())
101+
)
102+
}
103+
(for_target, arg) => format!(
104+
"{}.{batch_method_name}({} for {} in {})",
105+
set.id,
106+
checker.locator().slice(arg),
107+
checker.locator().slice(for_target),
108+
checker.locator().slice(for_stmt.iter.as_ref())
109+
),
110+
};
111+
112+
checker.diagnostics.push(
113+
Diagnostic::new(
114+
ForLoopSetMutations {
115+
method_name,
116+
batch_method_name,
117+
},
118+
for_stmt.range,
119+
)
120+
.with_fix(Fix::safe_edit(Edit::range_replacement(
121+
content,
122+
for_stmt.range,
123+
))),
124+
);
125+
}

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

+2
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
pub(crate) use bit_count::*;
22
pub(crate) use check_and_remove_from_set::*;
33
pub(crate) use delete_full_slice::*;
4+
pub(crate) use for_loop_set_mutations::*;
45
pub(crate) use hashlib_digest_hex::*;
56
pub(crate) use if_expr_min_max::*;
67
pub(crate) use implicit_cwd::*;
@@ -25,6 +26,7 @@ pub(crate) use verbose_decimal_constructor::*;
2526
mod bit_count;
2627
mod check_and_remove_from_set;
2728
mod delete_full_slice;
29+
mod for_loop_set_mutations;
2830
mod hashlib_digest_hex;
2931
mod if_expr_min_max;
3032
mod implicit_cwd;

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

+161
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,161 @@
1+
---
2+
source: crates/ruff_linter/src/rules/refurb/mod.rs
3+
---
4+
FURB142.py:5:1: FURB142 [*] Use of set.add in a for loop
5+
|
6+
3 | s = set()
7+
4 |
8+
5 | / for x in [1, 2, 3]:
9+
6 | | s.add(x)
10+
| |____________^ FURB142
11+
7 |
12+
8 | for x in {1, 2, 3}:
13+
|
14+
= help: Replace with `update`
15+
16+
ℹ Safe fix
17+
2 2 |
18+
3 3 | s = set()
19+
4 4 |
20+
5 |-for x in [1, 2, 3]:
21+
6 |- s.add(x)
22+
5 |+s.update([1, 2, 3])
23+
7 6 |
24+
8 7 | for x in {1, 2, 3}:
25+
9 8 | s.add(x)
26+
27+
FURB142.py:8:1: FURB142 [*] Use of set.add in a for loop
28+
|
29+
6 | s.add(x)
30+
7 |
31+
8 | / for x in {1, 2, 3}:
32+
9 | | s.add(x)
33+
| |____________^ FURB142
34+
10 |
35+
11 | for x in (1, 2, 3):
36+
|
37+
= help: Replace with `update`
38+
39+
ℹ Safe fix
40+
5 5 | for x in [1, 2, 3]:
41+
6 6 | s.add(x)
42+
7 7 |
43+
8 |-for x in {1, 2, 3}:
44+
9 |- s.add(x)
45+
8 |+s.update({1, 2, 3})
46+
10 9 |
47+
11 10 | for x in (1, 2, 3):
48+
12 11 | s.add(x)
49+
50+
FURB142.py:11:1: FURB142 [*] Use of set.add in a for loop
51+
|
52+
9 | s.add(x)
53+
10 |
54+
11 | / for x in (1, 2, 3):
55+
12 | | s.add(x)
56+
| |____________^ FURB142
57+
13 |
58+
14 | for x in (1, 2, 3):
59+
|
60+
= help: Replace with `update`
61+
62+
ℹ Safe fix
63+
8 8 | for x in {1, 2, 3}:
64+
9 9 | s.add(x)
65+
10 10 |
66+
11 |-for x in (1, 2, 3):
67+
12 |- s.add(x)
68+
11 |+s.update((1, 2, 3))
69+
13 12 |
70+
14 13 | for x in (1, 2, 3):
71+
15 14 | s.discard(x)
72+
73+
FURB142.py:14:1: FURB142 [*] Use of set.discard in a for loop
74+
|
75+
12 | s.add(x)
76+
13 |
77+
14 | / for x in (1, 2, 3):
78+
15 | | s.discard(x)
79+
| |________________^ FURB142
80+
16 |
81+
17 | for x in (1, 2, 3):
82+
|
83+
= help: Replace with `difference_update`
84+
85+
ℹ Safe fix
86+
11 11 | for x in (1, 2, 3):
87+
12 12 | s.add(x)
88+
13 13 |
89+
14 |-for x in (1, 2, 3):
90+
15 |- s.discard(x)
91+
14 |+s.difference_update((1, 2, 3))
92+
16 15 |
93+
17 16 | for x in (1, 2, 3):
94+
18 17 | s.add(x + 1)
95+
96+
FURB142.py:17:1: FURB142 [*] Use of set.add in a for loop
97+
|
98+
15 | s.discard(x)
99+
16 |
100+
17 | / for x in (1, 2, 3):
101+
18 | | s.add(x + 1)
102+
| |________________^ FURB142
103+
19 |
104+
20 | for x, y in ((1, 2), (3, 4)):
105+
|
106+
= help: Replace with `update`
107+
108+
ℹ Safe fix
109+
14 14 | for x in (1, 2, 3):
110+
15 15 | s.discard(x)
111+
16 16 |
112+
17 |-for x in (1, 2, 3):
113+
18 |- s.add(x + 1)
114+
17 |+s.update(x + 1 for x in (1, 2, 3))
115+
19 18 |
116+
20 19 | for x, y in ((1, 2), (3, 4)):
117+
21 20 | s.add((x, y))
118+
119+
FURB142.py:20:1: FURB142 [*] Use of set.add in a for loop
120+
|
121+
18 | s.add(x + 1)
122+
19 |
123+
20 | / for x, y in ((1, 2), (3, 4)):
124+
21 | | s.add((x, y))
125+
| |_________________^ FURB142
126+
22 |
127+
23 | num = 123
128+
|
129+
= help: Replace with `update`
130+
131+
ℹ Safe fix
132+
17 17 | for x in (1, 2, 3):
133+
18 18 | s.add(x + 1)
134+
19 19 |
135+
20 |-for x, y in ((1, 2), (3, 4)):
136+
21 |- s.add((x, y))
137+
20 |+s.update((x, y) for x, y in ((1, 2), (3, 4)))
138+
22 21 |
139+
23 22 | num = 123
140+
24 23 |
141+
142+
FURB142.py:25:1: FURB142 [*] Use of set.add in a for loop
143+
|
144+
23 | num = 123
145+
24 |
146+
25 | / for x in (1, 2, 3):
147+
26 | | s.add(num)
148+
| |______________^ FURB142
149+
|
150+
= help: Replace with `update`
151+
152+
ℹ Safe fix
153+
22 22 |
154+
23 23 | num = 123
155+
24 24 |
156+
25 |-for x in (1, 2, 3):
157+
26 |- s.add(num)
158+
25 |+s.update(num for x in (1, 2, 3))
159+
27 26 |
160+
28 27 |
161+
29 28 | # False negative

‎ruff.schema.json

+1
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)