Add use-set-for-membership check

ChangeLog

@@ -10,6 +10,11 @@ Release date: TBA
 ..
   Put new features here and also in 'doc/whatsnew/2.11.rst'
 
+* Added new extension ``SetMembershipChecker`` with ``use-set-for-membership`` check:
+  Emitted when using an in-place defined ``list`` or ``tuple`` to do a membership test. ``sets`` are better optimized for that.
+
+  Closes #4776
+
 
 * Added ``py-version`` config key (if ``[MASTER]`` section). Used for version dependant checks.
   Will default to whatever Python version pylint is executed with.

doc/whatsnew/2.11.rst

@@ -20,6 +20,13 @@ New checkers
 Extensions
 ==========
 
+
+* Added new extension ``SetMembershipChecker`` with ``use-set-for-membership`` check:
+  Emitted when using an in-place defined ``list`` or ``tuple`` to do a membership test. ``sets`` are better optimized for that.
+
+  Closes #4776
+
+
 * ``CodeStyleChecker``
 
   * Added ``consider-using-assignment-expr``: Emitted when an assignment is directly followed by an if statement

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))

tests/functional/ext/set_membership/set_membership_use_set_membership.py

@@ -0,0 +1,40 @@
+# pylint: disable=invalid-name,missing-docstring,pointless-statement,unnecessary-comprehension,undefined-variable
+
+x = 1
+var = frozenset({1, 2, 3})
+
+x in var
+x in {1, 2, 3}
+x in (1, 2, 3)  # [use-set-for-membership]
+x in [1, 2, 3]  # [use-set-for-membership]
+
+if x in var:
+    pass
+if x in {1, 2, 3}:
+    pass
+if x in (1, 2, 3):  # [use-set-for-membership]
+    pass
+if x in [1, 2, 3]:  # [use-set-for-membership]
+    pass
+
+42 if x in [1, 2, 3] else None  # [use-set-for-membership]
+assert x in [1, 2, 3]  # [use-set-for-membership]
+(x for x in var if x in [1, 2, 3])  # [use-set-for-membership]
+while x in [1, 2, 3]:  # [use-set-for-membership]
+    break
+
+# Stacked operators, rightmost pair is evaluated first
+# Doesn't make much sense in practice since `in` will only return `bool`
+True == x in [1, 2, 3]  # [use-set-for-membership]  # noqa: E712
+1 >= x in [1, 2, 3]  # [use-set-for-membership]  # noqa: E712
+
+
+# Test hashable heuristic
+x in (1, "Hello World", False, None)  # [use-set-for-membership]
+x in (1, [])  # List is not hashable
+
+if some_var:
+    var2 = 2
+else:
+    var2 = []
+x in (1, var2)  # var2 can be a list

tests/functional/ext/set_membership/set_membership_use_set_membership.rc

@@ -0,0 +1,2 @@
+[MASTER]
+load-plugins=pylint.extensions.set_membership

tests/functional/ext/set_membership/set_membership_use_set_membership.txt

@@ -0,0 +1,11 @@
+use-set-for-membership:8:5::Consider using set for membership test:HIGH
+use-set-for-membership:9:5::Consider using set for membership test:HIGH
+use-set-for-membership:15:8::Consider using set for membership test:HIGH
+use-set-for-membership:17:8::Consider using set for membership test:HIGH
+use-set-for-membership:20:11::Consider using set for membership test:HIGH
+use-set-for-membership:21:12::Consider using set for membership test:HIGH
+use-set-for-membership:22:24::Consider using set for membership test:HIGH
+use-set-for-membership:23:11::Consider using set for membership test:HIGH
+use-set-for-membership:28:13::Consider using set for membership test:HIGH
+use-set-for-membership:29:10::Consider using set for membership test:HIGH
+use-set-for-membership:33:5::Consider using set for membership test:HIGH