2021-12-27 20:30:34 +00:00
|
|
|
from typing import TYPE_CHECKING
|
|
|
|
|
2021-08-30 16:38:31 +00:00
|
|
|
from astroid import nodes
|
|
|
|
|
|
|
|
from pylint.checkers import BaseChecker
|
|
|
|
from pylint.checkers.utils import check_messages
|
|
|
|
from pylint.interfaces import IAstroidChecker
|
2021-12-27 20:30:34 +00:00
|
|
|
|
|
|
|
if TYPE_CHECKING:
|
|
|
|
from pylint.lint import PyLinter
|
2021-08-30 16:38:31 +00:00
|
|
|
|
|
|
|
|
|
|
|
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``.",
|
|
|
|
),
|
|
|
|
}
|
|
|
|
|
2021-12-27 20:30:34 +00:00
|
|
|
def __init__(self, linter: "PyLinter") -> None:
|
2021-08-30 16:38:31 +00:00
|
|
|
"""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)
|
|
|
|
|
|
|
|
|
2021-12-27 20:30:34 +00:00
|
|
|
def register(linter: "PyLinter") -> None:
|
2021-08-30 16:38:31 +00:00
|
|
|
linter.register_checker(SetMembershipChecker(linter))
|