diff --git a/docs/changelog.md b/docs/changelog.md index d3e1bcaa..f9a24e6d 100644 --- 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 index 0fc16edc..cca9ea08 100644 --- 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 index b0ee16bd..f18c390c 100644 --- 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: