Repository: quora/pyanalyze
Pull request #297 (2021-12-03 00:09:23)

More Protocol support
This supports types that are marked as Protocols in typeshed. We're still missing support for runtime Protocols.


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index 2a175f9..704bc74 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Support Protocols for runtime types that are also defined in stubs (#297)
 - 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 e251b77..2e5fd2a 100644
--- a/pyanalyze/checker.py
+++ b/pyanalyze/checker.py
@@ -12,7 +12,7 @@ from .safe import is_typing_name
 from .type_object import TypeObject
 
 from dataclasses import dataclass, field
-from typing import Set, Union, Dict
+from typing import Iterable, Set, Union, Dict
 
 
 @dataclass
@@ -45,21 +45,37 @@ class Checker:
     def _build_type_object(self, typ: Union[type, super, str]) -> TypeObject:
         if isinstance(typ, str):
             # Synthetic type
-            base_values = self.arg_spec_cache.ts_finder.get_bases_recursively(typ)
-            bases = set(
-                base.typ for base in base_values if isinstance(base, TypedValue)
-            )
+            bases = self._get_typeshed_bases(typ)
             is_protocol = any(is_typing_name(base, "Protocol") for base in bases)
             if is_protocol:
-                protocol_members = set(
-                    itertools.chain.from_iterable(
-                        self.arg_spec_cache.ts_finder.get_all_attributes(base)
-                        for base in bases
-                    )
-                )
+                protocol_members = self._get_protocol_members(bases)
             else:
                 protocol_members = set()
             return TypeObject(
                 typ, bases, is_protocol=is_protocol, protocol_members=protocol_members
             )
-        return TypeObject(typ, self.get_additional_bases(typ))
+        elif isinstance(typ, super):
+            return TypeObject(typ, self.get_additional_bases(typ))
+        else:
+            additional_bases = self.get_additional_bases(typ)
+            is_protocol = self.arg_spec_cache.ts_finder.is_protocol(typ)
+            if is_protocol:
+                bases = self._get_typeshed_bases(typ)
+                return TypeObject(
+                    typ,
+                    additional_bases,
+                    is_protocol=True,
+                    protocol_members=self._get_protocol_members(bases),
+                )
+            return TypeObject(typ, additional_bases)
+
+    def _get_typeshed_bases(self, typ: Union[type, str]) -> Set[Union[type, str]]:
+        base_values = self.arg_spec_cache.ts_finder.get_bases_recursively(typ)
+        return set(base.typ for base in base_values if isinstance(base, TypedValue))
+
+    def _get_protocol_members(self, bases: Iterable[Union[type, str]]) -> Set[str]:
+        return set(
+            itertools.chain.from_iterable(
+                self.arg_spec_cache.ts_finder.get_all_attributes(base) for base in bases
+            )
+        )
diff --git a/pyanalyze/typeshed.py b/pyanalyze/typeshed.py
index 9f33b1a..80d8dcd 100644
--- a/pyanalyze/typeshed.py
+++ b/pyanalyze/typeshed.py
@@ -187,11 +187,34 @@ class TypeshedFinder:
                     return None
             else:
                 fq_name = val.typ
+                if fq_name == "collections.abc.Set":
+                    return [GenericValue(Collection, (TypeVarValue(T_co),))]
+                elif fq_name == "contextlib.AbstractContextManager":
+                    return [GenericValue(Generic, (TypeVarValue(T_co),))]
+                elif fq_name in ("typing.Callable", "collections.abc.Callable"):
+                    return None
+                elif is_typing_name(fq_name, "TypedDict"):
+                    return [
+                        GenericValue(
+                            MutableMapping, [TypedValue(str), TypedValue(object)]
+                        )
+                    ]
             return self.get_bases_for_fq_name(fq_name)
         return None
 
-    def get_bases_recursively(self, fq_name: str) -> List[Value]:
-        stack = [TypedValue(fq_name)]
+    def is_protocol(self, typ: type) -> bool:
+        """Return whether this type is marked as a Protocol in the stubs."""
+        fq_name = self._get_fq_name(typ)
+        if fq_name is None:
+            return False
+        bases = self.get_bases_recursively(fq_name)
+        return any(
+            isinstance(base, TypedValue) and is_typing_name(base.typ, "Protocol")
+            for base in bases
+        )
+
+    def get_bases_recursively(self, typ: Union[type, str]) -> List[Value]:
+        stack = [TypedValue(typ)]
         seen = set()
         bases = []
         # TODO return MRO order


Test patch:
diff --git a/pyanalyze/test_type_object.py b/pyanalyze/test_type_object.py
index b4bd689..5a0adb8 100644
--- a/pyanalyze/test_type_object.py
+++ b/pyanalyze/test_type_object.py
@@ -153,3 +153,48 @@ class TestSyntheticType(TestNameCheckVisitorBase):
             cgi.parse(environ=Good())
             cgi.parse(environ=Bad())  # E: incompatible_argument
             cgi.parse(environ=1)  # E: incompatible_argument
+
+    @assert_passes()
+    def test_iterable(self):
+        from typing import Iterable, Iterator
+
+        class Bad:
+            def __iter__(self, some, random, args):
+                pass
+
+        class Good:
+            def __iter__(self) -> Iterator[int]:
+                raise NotImplementedError
+
+        class BadType:
+            def __iter__(self) -> Iterator[str]:
+                raise NotImplementedError
+
+        def want_iter_int(f: Iterable[int]) -> None:
+            pass
+
+        def capybara():
+            want_iter_int(Bad())  # E: incompatible_argument
+            want_iter_int(Good())
+            want_iter_int(BadType())  # E: incompatible_argument
+
+    @assert_passes()
+    def test_container(self):
+        from typing import Container, Any
+
+        class Good:
+            def __contains__(self, whatever: object) -> bool:
+                return False
+
+        class Bad:
+            def __contains__(self, too, many, arguments) -> bool:
+                return True
+
+        def want_container(c: Container[Any]) -> None:
+            pass
+
+        def capybara() -> None:
+            want_container(Bad())  # E: incompatible_argument
+            want_container(Good())
+            want_container([1])
+            want_container(1)  # E: incompatible_argument