set_membership.py 1.8 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152
  1. # Licensed under the GPL: https://www.gnu.org/licenses/old-licenses/gpl-2.0.html
  2. # For details: https://github.com/PyCQA/pylint/blob/main/LICENSE
  3. # Copyright (c) https://github.com/PyCQA/pylint/blob/main/CONTRIBUTORS.txt
  4. from __future__ import annotations
  5. from typing import TYPE_CHECKING
  6. from astroid import nodes
  7. from pylint.checkers import BaseChecker
  8. from pylint.checkers.utils import only_required_for_messages
  9. if TYPE_CHECKING:
  10. from pylint.lint import PyLinter
  11. class SetMembershipChecker(BaseChecker):
  12. name = "set_membership"
  13. msgs = {
  14. "R6201": (
  15. "Consider using set for membership test",
  16. "use-set-for-membership",
  17. "Membership tests are more efficient when performed on "
  18. "a lookup optimized datatype like ``sets``.",
  19. ),
  20. }
  21. def __init__(self, linter: PyLinter) -> None:
  22. """Initialize checker instance."""
  23. super().__init__(linter=linter)
  24. @only_required_for_messages("use-set-for-membership")
  25. def visit_compare(self, node: nodes.Compare) -> None:
  26. for op, comparator in node.ops:
  27. if op == "in":
  28. self._check_in_comparison(comparator)
  29. def _check_in_comparison(self, comparator: nodes.NodeNG) -> None:
  30. """Checks for membership comparisons with in-place container objects."""
  31. if not isinstance(comparator, nodes.BaseContainer) or isinstance(
  32. comparator, nodes.Set
  33. ):
  34. return
  35. # Heuristic - We need to be sure all items in set are hashable
  36. if all(isinstance(item, nodes.Const) for item in comparator.elts):
  37. self.add_message("use-set-for-membership", node=comparator)
  38. def register(linter: PyLinter) -> None:
  39. linter.register_checker(SetMembershipChecker(linter))