Support "in" on objects with __getitem__ (#564)

Fixes #475
quora · Nov 7, 2022 · 825b380 · 825b380
1 parent a35629a
commit 825b380
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 0 deletions.
diff --git a/docs/changelog.md b/docs/changelog.md
@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Support `in` on objects with only `__getitem__` (#564)
 - Add support for `except*` (PEP 654) (#562)
 - Add type inference support for more constructs in `except` and `except*` (#562)
 

diff --git a/pyanalyze/name_check_visitor.py b/pyanalyze/name_check_visitor.py
@@ -3175,6 +3175,32 @@ def _visit_binop_no_mvv(
             type(op)
         ]
         if rmethod is None:
+            # "in" falls back to __getitem__ if __contains__ is not defined
+            if method == "__contains__":
+                with self.catch_errors() as contains_errors:
+                    contains_result = self._check_dunder_call(
+                        source_node,
+                        left_composite,
+                        method,
+                        [right_composite],
+                        allow_call=allow_call,
+                    )
+                if not contains_errors:
+                    return contains_result
+
+                with self.catch_errors() as getitem_errors:
+                    self._check_dunder_call(
+                        source_node,
+                        left_composite,
+                        "__getitem__",
+                        [right_composite],
+                        allow_call=allow_call,
+                    )
+                if not getitem_errors:
+                    return TypedValue(bool)  # Always returns a bool
+                self.show_caught_errors(contains_errors)
+                return TypedValue(bool)
+
             return self._check_dunder_call(
                 source_node,
                 left_composite,

diff --git a/pyanalyze/test_operations.py b/pyanalyze/test_operations.py
@@ -231,6 +231,17 @@ def comparison(i: int, f: float, s: str, os: Optional[str]):
             s > None
             s > os
 
+    @assert_passes()
+    def test_contains(self):
+        class OnlyGetitem:
+            def __getitem__(self, x: int) -> int:
+                return x
+
+        def capybara(x: int, ogi: OnlyGetitem):
+            1 in x  # E: unsupported_operation
+            1 in ogi
+            "x" in ogi  # E: unsupported_operation
+
     @assert_passes()
     def test_failing_eq(self):
         class FlakyCapybara: