1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
|
from dataclasses import dataclass
from mypy.nodes import Block, CallExpr, ExpressionStmt, ForStmt, MemberExpr, NameExpr, Var
from refurb.checks.common import unmangle_name
from refurb.error import Error
@dataclass
class ErrorInfo(Error):
"""
When you want to add/remove a bunch of items to/from a set, don't use a for
loop, call the appropriate method on the set itself.
Bad:
```
sentence = "hello world"
vowels = "aeiou"
letters = set(sentence)
for vowel in vowels:
letters.discard(vowel)
```
Good:
```
sentence = "hello world"
vowels = "aeiou"
letters = set(sentence)
letters.difference_update(vowels)
```
"""
name = "no-set-for-loop"
code = 142
categories = ("builtin",)
def check(node: ForStmt, errors: list[Error]) -> None:
match node:
case ForStmt(
index=NameExpr(name=for_name),
body=Block(
body=[
ExpressionStmt(
expr=CallExpr(
callee=MemberExpr(
expr=NameExpr(node=Var(type=ty)) as set_name,
name=("add" | "discard") as name,
),
args=[arg],
)
)
]
),
) if str(ty).startswith("builtins.set[") and set_name.name != for_name:
new_func = "update" if name == "add" else "difference_update"
if isinstance(arg, NameExpr):
expr = unmangle_name(arg.name)
new_expr = "y"
if unmangle_name(for_name) != expr:
return
else:
expr = "..."
new_expr = "... for x in y"
errors.append(
ErrorInfo.from_node(
node,
f"Replace `for x in y: s.{name}({expr})` with `s.{new_func}({new_expr})`", # noqa: E501
)
)
|
