blob: c63e268b3cab5bff6f3edad028f101359ed237b0 (
plain)
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
|
from astroid import nodes
from pylint.checkers import BaseChecker
from pylint.checkers.utils import check_messages
from pylint.interfaces import IAstroidChecker
from pylint.lint import PyLinter
class SetMembershipChecker(BaseChecker):
__implements__ = (IAstroidChecker,)
name = "set_membership"
priority = -1
msgs = {
"R6201": (
"Consider using set for membership test",
"use-set-for-membership",
"Membership tests are more efficient when performed on "
"a lookup optimized datatype like ``sets``.",
),
}
def __init__(self, linter: PyLinter) -> None:
"""Initialize checker instance."""
super().__init__(linter=linter)
@check_messages("use-set-for-membership")
def visit_compare(self, node: nodes.Compare) -> None:
for op, comparator in node.ops:
if op == "in":
self._check_in_comparison(comparator)
def _check_in_comparison(self, comparator: nodes.NodeNG) -> None:
"""Checks for membership comparisons with in-place container objects."""
if not isinstance(comparator, nodes.BaseContainer) or isinstance(
comparator, nodes.Set
):
return
# Heuristic - We need to be sure all items in set are hashable
if all(isinstance(item, nodes.Const) for item in comparator.elts):
self.add_message("use-set-for-membership", node=comparator)
def register(linter: PyLinter) -> None:
linter.register_checker(SetMembershipChecker(linter))
|