Repository: quora/pyanalyze
Pull request #588 (2023-01-16 20:03:03)

Improve `in` handling


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index 6244b16..7a18185 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Support `in` on objects with only `__iter__` (#588)
 - Do not call `.mro()` method on non-types (#587)
 - Add `class_attribute_transformers` hook (#585)
 - Support for PEP 702 (`@typing.deprecated`) (#578)
diff --git a/pyanalyze/functions.py b/pyanalyze/functions.py
index b746e2f..9edf6b1 100644
--- a/pyanalyze/functions.py
+++ b/pyanalyze/functions.py
@@ -23,7 +23,6 @@ from .node_visitor import ErrorContext
 from .options import Options, PyObjectSequenceOption
 from .signature import ParameterKind, Signature, SigParameter
 from .stacked_scopes import Composite
-from .typevar import resolve_bounds_map
 from .value import (
     AnySource,
     AnyValue,
@@ -33,6 +32,7 @@ from .value import (
     GenericValue,
     get_tv_map,
     KnownValue,
+    is_iterable,
     make_coro_type,
     SubclassValue,
     TypedValue,
@@ -49,7 +49,6 @@ IMPLICIT_CLASSMETHODS = ("__init_subclass__", "__new__")
 YieldT = TypeVar("YieldT")
 SendT = TypeVar("SendT")
 ReturnT = TypeVar("ReturnT")
-IterableValue = GenericValue(collections.abc.Iterable, [TypeVarValue(YieldT)])
 GeneratorValue = GenericValue(
     collections.abc.Generator,
     [TypeVarValue(YieldT), TypeVarValue(SendT), TypeVarValue(ReturnT)],
@@ -93,11 +92,10 @@ class FunctionInfo:
     def get_generator_yield_type(self, ctx: CanAssignContext) -> Value:
         if self.return_annotation is None:
             return AnyValue(AnySource.unannotated)
-        can_assign = IterableValue.can_assign(self.return_annotation, ctx)
-        if isinstance(can_assign, CanAssignError):
+        iterable_val = is_iterable(self.return_annotation, ctx)
+        if isinstance(iterable_val, CanAssignError):
             return AnyValue(AnySource.error)
-        tv_map, _ = resolve_bounds_map(can_assign, ctx)
-        return tv_map.get(YieldT, AnyValue(AnySource.generic_argument))
+        return iterable_val
 
     def get_generator_send_type(self, ctx: CanAssignContext) -> Value:
         if self.return_annotation is None:
@@ -107,8 +105,8 @@ class FunctionInfo:
             return tv_map.get(SendT, AnyValue(AnySource.generic_argument))
         # If the return annotation is a non-Generator Iterable, assume the send
         # type is None.
-        can_assign = IterableValue.can_assign(self.return_annotation, ctx)
-        if isinstance(can_assign, CanAssignError):
+        iterable_val = is_iterable(self.return_annotation, ctx)
+        if isinstance(iterable_val, CanAssignError):
             return AnyValue(AnySource.error)
         return KnownValue(None)
 
@@ -120,8 +118,8 @@ class FunctionInfo:
             return tv_map.get(ReturnT, AnyValue(AnySource.generic_argument))
         # If the return annotation is a non-Generator Iterable, assume the return
         # type is None.
-        can_assign = IterableValue.can_assign(self.return_annotation, ctx)
-        if isinstance(can_assign, CanAssignError):
+        iterable_val = is_iterable(self.return_annotation, ctx)
+        if isinstance(iterable_val, CanAssignError):
             return AnyValue(AnySource.error)
         return KnownValue(None)
 
diff --git a/pyanalyze/implementation.py b/pyanalyze/implementation.py
index 400266d..7f6e1b7 100644
--- a/pyanalyze/implementation.py
+++ b/pyanalyze/implementation.py
@@ -375,7 +375,7 @@ def _sequence_getitem_impl(ctx: CallContext, typ: type) -> ImplReturn:
             return AnyValue(AnySource.error)  # shouldn't happen
         key = replace_known_sequence_value(key)
         if not TypedValue(slice).is_assignable(key, ctx.visitor):
-            key = ctx.visitor._check_dunder_call(
+            key, _ = ctx.visitor._check_dunder_call(
                 ctx.ast_for_arg("obj"), Composite(key), "__index__", [], allow_call=True
             )
 
diff --git a/pyanalyze/name_check_visitor.py b/pyanalyze/name_check_visitor.py
index d5f178e..13d21d3 100644
--- a/pyanalyze/name_check_visitor.py
+++ b/pyanalyze/name_check_visitor.py
@@ -79,7 +79,6 @@ from .functions import (
     FunctionResult,
     GeneratorValue,
     IMPLICIT_CLASSMETHODS,
-    IterableValue,
     ReturnT,
     SendT,
     YieldT,
@@ -173,6 +172,7 @@ from .value import (
     GenericBases,
     GenericValue,
     get_tv_map,
+    is_iterable,
     is_union,
     KnownValue,
     kv_pairs_from_mapping,
@@ -3234,7 +3234,7 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
         else:
             operand = self.composite_from_node(node.operand)
             _, method = UNARY_OPERATION_TO_DESCRIPTION_AND_METHOD[type(node.op)]
-            val = self._check_dunder_call(node, operand, method, [], allow_call=True)
+            val, _ = self._check_dunder_call(node, operand, method, [], allow_call=True)
             return val
 
     def visit_BinOp(self, node: ast.BinOp) -> Value:
@@ -3299,7 +3299,10 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
         if is_inplace:
             assert imethod is not None, f"no inplace method available for {op}"
             with self.catch_errors() as inplace_errors:
-                inplace_result = self._check_dunder_call(
+                # Not _check_dunder_call_or_catch because if the call doesn't
+                # typecheck it normally returns NotImplemented and we try the
+                # non-inplace method next.
+                inplace_result, _ = self._check_dunder_call(
                     source_node,
                     left_composite,
                     imethod,
@@ -3335,42 +3338,55 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
             type(op)
         ]
         if rmethod is None:
-            # "in" falls back to __getitem__ if __contains__ is not defined
+            # "in" falls back to __iter__ and then 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:
+                contains_result_or_errors = self._check_dunder_call_or_catch(
+                    source_node,
+                    left_composite,
+                    method,
+                    [right_composite],
+                    allow_call=allow_call,
+                )
+                if isinstance(contains_result_or_errors, Value):
+                    return contains_result_or_errors
+
+                iterable_type = is_iterable(left, self)
+                if isinstance(iterable_type, Value):
+                    can_assign = iterable_type.can_assign(right, self)
+                    if isinstance(can_assign, CanAssignError):
+                        self._show_error_if_checking(
+                            source_node,
+                            "Unsupported operand for 'in'",
+                            ErrorCode.incompatible_argument,
+                            detail=str(can_assign),
+                        )
+                        return TypedValue(bool)
+                    else:
+                        return TypedValue(bool)
+
+                getitem_result = self._check_dunder_call_or_catch(
+                    source_node,
+                    left_composite,
+                    "__getitem__",
+                    [right_composite],
+                    allow_call=allow_call,
+                )
+                if isinstance(getitem_result, Value):
                     return TypedValue(bool)  # Always returns a bool
-                self.show_caught_errors(contains_errors)
+                self.show_caught_errors(contains_result_or_errors)
                 return TypedValue(bool)
 
-            return self._check_dunder_call(
+            result, _ = self._check_dunder_call(
                 source_node,
                 left_composite,
                 method,
                 [right_composite],
                 allow_call=allow_call,
             )
+            return result
 
         with self.catch_errors() as left_errors:
-            left_result = self._check_dunder_call(
+            left_result, _ = self._check_dunder_call(
                 source_node,
                 left_composite,
                 method,
@@ -3379,7 +3395,7 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
             )
 
         with self.catch_errors() as right_errors:
-            right_result = self._check_dunder_call(
+            right_result, _ = self._check_dunder_call(
                 source_node,
                 right_composite,
                 rmethod,
@@ -3458,7 +3474,8 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
     def unpack_awaitable(self, composite: Composite, node: ast.AST) -> Value:
         tv_map = get_tv_map(AwaitableValue, composite.value, self)
         if isinstance(tv_map, CanAssignError):
-            return self._check_dunder_call(node, composite, "__await__", [])
+            result, _ = self._check_dunder_call(node, composite, "__await__", [])
+            return result
         else:
             return tv_map.get(T, AnyValue(AnySource.generic_argument))
 
@@ -3473,8 +3490,8 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
                     ReturnT: can_assign.get(T, AnyValue(AnySource.generic_argument))
                 }
             else:
-                can_assign = get_tv_map(IterableValue, value, self)
-                if isinstance(can_assign, CanAssignError):
+                iterable_type = is_iterable(value, self)
+                if isinstance(iterable_type, CanAssignError):
                     self._show_error_if_checking(
                         node,
                         f"Cannot use {value} in yield from",
@@ -3483,9 +3500,7 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
                     )
                     tv_map = {ReturnT: AnyValue(AnySource.error)}
                 else:
-                    tv_map = {
-                        YieldT: can_assign.get(T, AnyValue(AnySource.generic_argument))
-                    }
+                    tv_map = {YieldT: iterable_type}
 
         if self.current_function_info is not None:
             expected_yield = self.current_function_info.get_generator_yield_type(self)
@@ -3838,8 +3853,8 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
         """
         composite = self.composite_from_node(node)
         if is_async:
-            iterator = self._check_dunder_call(node, composite, "__aiter__", [])
-            anext = self._check_dunder_call(
+            iterator, _ = self._check_dunder_call(node, composite, "__aiter__", [])
+            anext, _ = self._check_dunder_call(
                 node, Composite(iterator, None, node), "__anext__", []
             )
             return self.unpack_awaitable(Composite(anext), node)
@@ -4468,9 +4483,10 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
                     return_value = local_value
             return return_value
         elif isinstance(node.ctx, ast.Del):
-            return self._check_dunder_call(
+            result, _ = self._check_dunder_call(
                 node.value, root_composite, "__delitem__", [index_composite]
             )
+            return result
         else:
             self.show_error(
                 node,
@@ -4495,6 +4511,37 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
             )
         return method_object
 
+    def _check_dunder_call_or_catch(
+        self,
+        node: ast.AST,
+        callee_composite: Composite,
+        method_name: str,
+        args: Iterable[Composite],
+        allow_call: bool = False,
+    ) -> Union[Value, List[node_visitor.Error]]:
+        """Use this for checking a dunder call that may fall back to another.
+
+        There are three cases:
+        - The dunder does not exist. We want to defer the error, in case the fallback
+          exists.
+        - The dunder exists and the call typechecks. We want to return its result.
+        - The dunder exists, but the call doesn't typecheck. We want to show the error
+          immediately and return Any.
+
+        """
+        with self.catch_errors() as errors:
+            result, exists = self._check_dunder_call(
+                node, callee_composite, method_name, args, allow_call=allow_call
+            )
+        if not errors:
+            return result
+        elif exists:
+            # Inplace method exists, but it doesn't accept these arguments
+            self.show_caught_errors(errors)
+            return result
+        else:
+            return errors
+
     def _check_dunder_call(
         self,
         node: ast.AST,
@@ -4502,21 +4549,27 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
         method_name: str,
         args: Iterable[Composite],
         allow_call: bool = False,
-    ) -> Value:
+    ) -> Tuple[Value, bool]:
         if isinstance(callee_composite.value, MultiValuedValue):
             composites = [
                 Composite(val, callee_composite.varname, callee_composite.node)
                 for val in callee_composite.value.vals
             ]
             with qcore.override(self, "in_union_decomposition", True):
-                values = [
+                values_and_exists = [
                     self._check_dunder_call_no_mvv(
                         node, composite, method_name, args, allow_call
                     )
                     for composite in composites
                 ]
-            return unite_and_simplify(
-                *values, limit=self.options.get_value_for(UnionSimplificationLimit)
+            values = [value for value, _ in values_and_exists]
+            # TODO: We should do something more complex when unions are involved.
+            exists = all(exists for _, exists in values_and_exists)
+            return (
+                unite_and_simplify(
+                    *values, limit=self.options.get_value_for(UnionSimplificationLimit)
+                ),
+                exists,
             )
         return self._check_dunder_call_no_mvv(
             node, callee_composite, method_name, args, allow_call
@@ -4529,14 +4582,14 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
         method_name: str,
         args: Iterable[Composite],
         allow_call: bool = False,
-    ) -> Value:
+    ) -> Tuple[Value, bool]:
         method_object = self._get_dunder(node, callee_composite.value, method_name)
         if method_object is UNINITIALIZED_VALUE:
-            return AnyValue(AnySource.error)
+            return AnyValue(AnySource.error), False
         return_value = self.check_call(
             node, method_object, [callee_composite, *args], allow_call=allow_call
         )
-        return return_value
+        return return_value, True
 
     def _get_composite(self, composite: Varname, node: ast.AST, value: Value) -> Value:
         local_value, _ = self.scopes.current_scope().get_local(
diff --git a/pyanalyze/node_visitor.py b/pyanalyze/node_visitor.py
index 4b717b0..75a25e4 100644
--- a/pyanalyze/node_visitor.py
+++ b/pyanalyze/node_visitor.py
@@ -44,6 +44,8 @@ from typing_extensions import NotRequired, Protocol, TypedDict
 
 from . import analysis_lib
 
+Error = Dict[str, Any]
+
 
 @dataclass(frozen=True)
 class _FakeNode:
@@ -541,12 +543,12 @@ class BaseNodeVisitor(ast.NodeVisitor):
             }
 
     @contextmanager
-    def catch_errors(self) -> Iterator[List[Dict[str, Any]]]:
+    def catch_errors(self) -> Iterator[List[Error]]:
         caught_errors = []
         with qcore.override(self, "caught_errors", caught_errors):
             yield caught_errors
 
-    def show_caught_errors(self, errors: Iterable[Dict[str, Any]]) -> None:
+    def show_caught_errors(self, errors: Iterable[Error]) -> None:
         for error in errors:
             self.show_error(**error)
 
diff --git a/pyanalyze/signature.py b/pyanalyze/signature.py
index 562bf4d..b957737 100644
--- a/pyanalyze/signature.py
+++ b/pyanalyze/signature.py
@@ -91,6 +91,7 @@ from .value import (
     ParameterTypeGuardExtension,
     ParamSpecArgsValue,
     ParamSpecKwargsValue,
+    is_iterable,
     replace_known_sequence_value,
     SequenceValue,
     stringify_object,
@@ -2541,8 +2542,6 @@ def make_bound_method(
         assert_never(argspec)
 
 
-T = TypeVar("T")
-IterableValue = GenericValue(collections.abc.Iterable, [TypeVarValue(T)])
 K = TypeVar("K")
 V = TypeVar("V")
 MappingValue = GenericValue(collections.abc.Mapping, [TypeVarValue(K), TypeVarValue(V)])
@@ -2573,10 +2572,11 @@ def can_assign_var_positional(
                 )
             return [can_assign]
 
-    tv_map = get_tv_map(IterableValue, args_annotation, ctx)
-    if isinstance(tv_map, CanAssignError):
-        return CanAssignError(f"{args_annotation} is not an iterable type", [tv_map])
-    iterable_arg = tv_map.get(T, AnyValue(AnySource.generic_argument))
+    iterable_arg = is_iterable(args_annotation, ctx)
+    if isinstance(iterable_arg, CanAssignError):
+        return CanAssignError(
+            f"{args_annotation} is not an iterable type", [iterable_arg]
+        )
     bounds_map = iterable_arg.can_assign(my_annotation, ctx)
     if isinstance(bounds_map, CanAssignError):
         return CanAssignError(
diff --git a/pyanalyze/value.py b/pyanalyze/value.py
index bb5d218..4224eb3 100644
--- a/pyanalyze/value.py
+++ b/pyanalyze/value.py
@@ -2344,9 +2344,9 @@ def concrete_values_from_iterable(
             is_nonempty = True
     elif value is NO_RETURN_VALUE:
         return NO_RETURN_VALUE
-    iter_tv_map = get_tv_map(IterableValue, value, ctx)
-    if not isinstance(iter_tv_map, CanAssignError):
-        val = iter_tv_map.get(T, AnyValue(AnySource.generic_argument))
+    iterable_type = is_iterable(value, ctx)
+    if isinstance(iterable_type, Value):
+        val = iterable_type
     else:
         getitem_tv_map = get_tv_map(GetItemProtoValue, value, ctx)
         if not isinstance(getitem_tv_map, CanAssignError):
@@ -2363,7 +2363,7 @@ def concrete_values_from_iterable(
         else:
             # We return the error from the __iter__ check because the __getitem__
             # check is more arcane.
-            return iter_tv_map
+            return iterable_type
     if is_nonempty:
         return annotate_value(val, [AlwaysPresentExtension()])
     return val
@@ -2543,11 +2543,18 @@ def unpack_values(
             if not isinstance(vals, CanAssignError):
                 return vals
 
+    iterable_type = is_iterable(value, ctx)
+    if isinstance(iterable_type, CanAssignError):
+        return iterable_type
+    return _create_unpacked_list(iterable_type, target_length, post_starred_length)
+
+
+def is_iterable(value: Value, ctx: CanAssignContext) -> Union[CanAssignError, Value]:
+    """Check whether a value is iterable."""
     tv_map = get_tv_map(IterableValue, value, ctx)
     if isinstance(tv_map, CanAssignError):
         return tv_map
-    iterable_type = tv_map.get(T, AnyValue(AnySource.generic_argument))
-    return _create_unpacked_list(iterable_type, target_length, post_starred_length)
+    return tv_map.get(T, AnyValue(AnySource.generic_argument))
 
 
 def _create_unpacked_list(


Test patch:
diff --git a/pyanalyze/test_operations.py b/pyanalyze/test_operations.py
index f18c390..80c0653 100644
--- a/pyanalyze/test_operations.py
+++ b/pyanalyze/test_operations.py
@@ -233,14 +233,24 @@ class TestCompare(TestNameCheckVisitorBase):
 
     @assert_passes()
     def test_contains(self):
+        from typing import Iterator, Iterable
+
         class OnlyGetitem:
             def __getitem__(self, x: int) -> int:
                 return x
 
-        def capybara(x: int, ogi: OnlyGetitem):
+        class OnlyIter:
+            def __iter__(self) -> Iterator[int]:
+                yield 1
+
+        def capybara(x: int, ogi: OnlyGetitem, oi: OnlyIter, it: Iterable[int]):
             1 in x  # E: unsupported_operation
             1 in ogi
-            "x" in ogi  # E: unsupported_operation
+            "x" in ogi  # E: incompatible_argument
+            1 in oi
+            "1" in oi  # E: incompatible_argument
+            1 in it
+            "1" in it  # E: incompatible_argument
 
     @assert_passes()
     def test_failing_eq(self):