Repository: quora/pyanalyze
Pull request #307 (2021-12-03 17:48:55)

fix recursion on Iterator


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index 5f64947..f490ccb 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -5,7 +5,7 @@
 - Fix false positive where `multiprocessing.Pool.map_async`
   was identified as an asynq method (#306)
 - Fix handling of nested classes (#305)
-- Support Protocols for runtime types that are also defined in stubs (#297)
+- Support Protocols for runtime types that are also defined in stubs (#297) (#307)
 - Better detect signatures of methods in stub files (#304)
 - Improve handling of positional-only arguments in stub files (#303)
 - Fix bug where pyanalyze incorrectly inferred that an attribute always exists (#302)
diff --git a/pyanalyze/checker.py b/pyanalyze/checker.py
index 2e5fd2a..73af99e 100644
--- a/pyanalyze/checker.py
+++ b/pyanalyze/checker.py
@@ -4,6 +4,10 @@ The checker maintains global state that is preserved across different modules.
 
 """
 import itertools
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from typing import Iterable, Iterator, List, Set, Tuple, Union, Dict
+
 from .value import TypedValue
 from .arg_spec import ArgSpecCache
 from .config import Config
@@ -11,9 +15,6 @@ from .reexport import ImplicitReexportTracker
 from .safe import is_typing_name
 from .type_object import TypeObject
 
-from dataclasses import dataclass, field
-from typing import Iterable, Set, Union, Dict
-
 
 @dataclass
 class Checker:
@@ -23,6 +24,9 @@ class Checker:
     type_object_cache: Dict[Union[type, super, str], TypeObject] = field(
         default_factory=dict, init=False, repr=False
     )
+    assumed_compatibilities: List[Tuple[TypeObject, TypeObject]] = field(
+        default_factory=list
+    )
 
     def __post_init__(self) -> None:
         self.arg_spec_cache = ArgSpecCache(self.config)
@@ -79,3 +83,19 @@ class Checker:
                 self.arg_spec_cache.ts_finder.get_all_attributes(base) for base in bases
             )
         )
+
+    def can_assume_compatibility(self, left: TypeObject, right: TypeObject) -> bool:
+        return (left, right) in self.assumed_compatibilities
+
+    @contextmanager
+    def assume_compatibility(
+        self, left: TypeObject, right: TypeObject
+    ) -> Iterator[None]:
+        """Context manager that notes that left and right can be assumed to be compatible."""
+        pair = (left, right)
+        self.assumed_compatibilities.append(pair)
+        try:
+            yield
+        finally:
+            new_pair = self.assumed_compatibilities.pop()
+            assert pair == new_pair
diff --git a/pyanalyze/name_check_visitor.py b/pyanalyze/name_check_visitor.py
index 88de71d..c3e5fc6 100644
--- a/pyanalyze/name_check_visitor.py
+++ b/pyanalyze/name_check_visitor.py
@@ -33,6 +33,7 @@ import traceback
 import types
 from typing import (
     ClassVar,
+    ContextManager,
     Iterator,
     Mapping,
     Iterable,
@@ -808,6 +809,14 @@ class NameCheckVisitor(
     def make_type_object(self, typ: Union[type, super]) -> TypeObject:
         return self.checker.make_type_object(typ)
 
+    def can_assume_compatibility(self, left: TypeObject, right: TypeObject) -> bool:
+        return self.checker.can_assume_compatibility(left, right)
+
+    def assume_compatibility(
+        self, left: TypeObject, right: TypeObject
+    ) -> ContextManager[None]:
+        return self.checker.assume_compatibility(left, right)
+
     # The type for typ should be type, but that leads Cython to reject calls that pass
     # an instance of ABCMeta.
     def get_generic_bases(
diff --git a/pyanalyze/type_object.py b/pyanalyze/type_object.py
index ff44c13..2497649 100644
--- a/pyanalyze/type_object.py
+++ b/pyanalyze/type_object.py
@@ -118,18 +118,22 @@ class TypeObject:
                 return CanAssignError(
                     f"Cannot assign super object {other} to protocol {self}"
                 )
-            tv_maps = []
-            for member in self.protocol_members:
-                expected = ctx.get_attribute_from_value(self_val, member)
-                actual = ctx.get_attribute_from_value(other_val, member)
-                if actual is UNINITIALIZED_VALUE:
-                    return CanAssignError(f"{other} has no attribute {member!r}")
-                tv_map = expected.can_assign(actual, ctx)
-                if isinstance(tv_map, CanAssignError):
-                    return CanAssignError(
-                        f"Value of protocol member {member!r} conflicts", [tv_map]
-                    )
-                tv_maps.append(tv_map)
+            # This is a guard against infinite recursion if the Protocol is recursive
+            if ctx.can_assume_compatibility(self, other):
+                return {}
+            with ctx.assume_compatibility(self, other):
+                tv_maps = []
+                for member in self.protocol_members:
+                    expected = ctx.get_attribute_from_value(self_val, member)
+                    actual = ctx.get_attribute_from_value(other_val, member)
+                    if actual is UNINITIALIZED_VALUE:
+                        return CanAssignError(f"{other} has no attribute {member!r}")
+                    tv_map = expected.can_assign(actual, ctx)
+                    if isinstance(tv_map, CanAssignError):
+                        return CanAssignError(
+                            f"Value of protocol member {member!r} conflicts", [tv_map]
+                        )
+                    tv_maps.append(tv_map)
             return unify_typevar_maps(tv_maps)
 
     def is_instance(self, obj: object) -> bool:
diff --git a/pyanalyze/value.py b/pyanalyze/value.py
index 626b4d6..48ed115 100644
--- a/pyanalyze/value.py
+++ b/pyanalyze/value.py
@@ -21,6 +21,7 @@ these subclasses and some related utilities.
 
 import collections.abc
 from collections import OrderedDict, defaultdict, deque
+from contextlib import contextmanager
 from dataclasses import dataclass, field, InitVar
 import enum
 import inspect
@@ -30,6 +31,7 @@ from typing import (
     Callable,
     Dict,
     Iterable,
+    Iterator,
     List,
     Mapping,
     Optional,
@@ -205,6 +207,21 @@ class CanAssignContext:
     def get_attribute_from_value(self, root_value: "Value", attribute: str) -> "Value":
         return UNINITIALIZED_VALUE
 
+    def can_assume_compatibility(
+        self,
+        left: "pyanalyze.type_object.TypeObject",
+        right: "pyanalyze.type_object.TypeObject",
+    ) -> bool:
+        return False
+
+    @contextmanager
+    def assume_compatibility(
+        self,
+        left: "pyanalyze.type_object.TypeObject",
+        right: "pyanalyze.type_object.TypeObject",
+    ) -> Iterator[None]:
+        yield
+
 
 @dataclass(frozen=True)
 class CanAssignError:


Test patch:
diff --git a/pyanalyze/test_type_object.py b/pyanalyze/test_type_object.py
index 5a0adb8..796c0a2 100644
--- a/pyanalyze/test_type_object.py
+++ b/pyanalyze/test_type_object.py
@@ -178,6 +178,23 @@ class TestSyntheticType(TestNameCheckVisitorBase):
             want_iter_int(Good())
             want_iter_int(BadType())  # E: incompatible_argument
 
+    @assert_passes()
+    def test_self_iterator(self):
+        from typing import Iterator
+
+        class MyIter:
+            def __iter__(self) -> "MyIter":
+                return self
+
+            def __next__(self) -> int:
+                return 42
+
+        def want_iter(it: Iterator[int]):
+            pass
+
+        def capybara():
+            want_iter(MyIter())
+
     @assert_passes()
     def test_container(self):
         from typing import Container, Any