Add `use-set-for-membership` check (#4841)

* Add use-set-for-membership check

* Create SetMembershipChecker extension

* Add heuristic to check items for hashability

Co-authored-by: Pierre Sassoulas <pierre.sassoulas@gmail.com>

1 files changed, 47 insertions, 0 deletions

diff --git a/pylint/extensions/set_membership.py b/pylint/extensions/set_membership.py
new file mode 100644
index 000000000..c63e268b3
--- /dev/null
+++ b/pylint/extensions/set_membership.py
@@ -0,0 +1,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))