diff --git a/pyproject.toml b/pyproject.toml index a2700a4..e6067eb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "refurb" -version = "1.3.0" +version = "1.4.0" description = "A tool for refurbish and modernize Python codebases" authors = ["dosisod"] license = "GPL-3.0-only" diff --git a/refurb/checks/builtin/set_discard.py b/refurb/checks/builtin/set_discard.py new file mode 100644 index 0000000..2c2a3ef --- /dev/null +++ b/refurb/checks/builtin/set_discard.py @@ -0,0 +1,69 @@ +from dataclasses import dataclass + +from mypy.nodes import ( + Block, + CallExpr, + ComparisonExpr, + ExpressionStmt, + IfStmt, + MemberExpr, + NameExpr, + Var, +) + +from refurb.error import Error + + +@dataclass +class ErrorInfo(Error): + """ + If you want to remove a value from a set regardless of whether it exists or + not, use the `discard()` method instead of `remove()`: + + Bad: + + ``` + nums = set((123, 456)) + + if 123 in nums: + nums.remove(123) + ``` + + Good: + + ``` + nums = set((123, 456)) + + nums.discard(123) + ``` + """ + + code = 132 + msg: str = "Replace `if x in s: s.remove(x)` with `s.discard(x)`" + + +def check(node: IfStmt, errors: list[Error]) -> None: + match node: + case IfStmt( + expr=[ComparisonExpr(operators=["in"], operands=[lhs, rhs])], + body=[ + Block( + body=[ + ExpressionStmt( + expr=CallExpr( + callee=MemberExpr( + expr=NameExpr(node=Var(type=ty)) as expr, + name="remove", + ), + args=[arg], + ) + ) + ] + ) + ], + ) if ( + str(lhs) == str(arg) + and str(rhs) == str(expr) + and str(ty).startswith("builtins.set[") + ): + errors.append(ErrorInfo(node.line, node.column)) diff --git a/test/data/err_132.py b/test/data/err_132.py new file mode 100644 index 0000000..25678e5 --- /dev/null +++ b/test/data/err_132.py @@ -0,0 +1,34 @@ +s = set() + +# these should match + +if "x" in s: + s.remove("x") + +# these should not + +if "x" in s: + s.remove("y") + +s.discard("x") + +s2 = set() + +if "x" in s: + s2.remove("x") + +if "x" in s: + s.remove("x") + print("removed item") + +class Container: + def remove(self, item) -> None: + return + + def __contains__(self, other) -> bool: + return True + +c = Container() + +if "x" in c: + c.remove("x") diff --git a/test/data/err_132.txt b/test/data/err_132.txt new file mode 100644 index 0000000..8d5d007 --- /dev/null +++ b/test/data/err_132.txt @@ -0,0 +1 @@ +test/data/err_132.py:5:1 [FURB132]: Replace `if x in s: s.remove(x)` with `s.discard(x)`