Repository: quora/pyanalyze
Pull request #295 (2021-12-02 20:55:43)

Basic Protocol support
Only synthetic, non-generic protocols.


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index 86b09f6..6c576e7 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Add support for stub-only Protocols (#295)
 - Basic support for stub-only types (#290)
 - Require `typing_inspect>=0.7.0` (#290)
 - Improve type checking of `raise` statements (#289)
diff --git a/pyanalyze/__init__.py b/pyanalyze/__init__.py
index f1034d0..8bc67b5 100644
--- a/pyanalyze/__init__.py
+++ b/pyanalyze/__init__.py
@@ -24,6 +24,7 @@ from . import safe
 from . import signature
 from . import stacked_scopes
 from . import test_config
+from . import type_object
 from . import typeshed
 from . import tests
 from . import value
diff --git a/pyanalyze/checker.py b/pyanalyze/checker.py
index fdd59be..e251b77 100644
--- a/pyanalyze/checker.py
+++ b/pyanalyze/checker.py
@@ -3,10 +3,12 @@
 The checker maintains global state that is preserved across different modules.
 
 """
+import itertools
 from .value import TypedValue
 from .arg_spec import ArgSpecCache
 from .config import Config
 from .reexport import ImplicitReexportTracker
+from .safe import is_typing_name
 from .type_object import TypeObject
 
 from dataclasses import dataclass, field
@@ -43,8 +45,21 @@ class Checker:
     def _build_type_object(self, typ: Union[type, super, str]) -> TypeObject:
         if isinstance(typ, str):
             # Synthetic type
-            bases = self.arg_spec_cache.ts_finder.get_bases_recursively(typ)
+            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)
+            )
+            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
+                    )
+                )
+            else:
+                protocol_members = set()
             return TypeObject(
-                typ, set(base.typ for base in bases if isinstance(base, TypedValue))
+                typ, bases, is_protocol=is_protocol, protocol_members=protocol_members
             )
         return TypeObject(typ, self.get_additional_bases(typ))
diff --git a/pyanalyze/name_check_visitor.py b/pyanalyze/name_check_visitor.py
index b878de2..ca7e500 100644
--- a/pyanalyze/name_check_visitor.py
+++ b/pyanalyze/name_check_visitor.py
@@ -3721,6 +3721,9 @@ class NameCheckVisitor(
             return unite_values(*values)
         return self._get_attribute_no_mvv(root_composite, attr, node)
 
+    def get_attribute_from_value(self, root_value: Value, attribute: str) -> Value:
+        return self.get_attribute(Composite(root_value), attribute)
+
     def _get_attribute_no_mvv(
         self, root_composite: Composite, attr: str, node: Optional[ast.AST] = None
     ) -> Value:
diff --git a/pyanalyze/type_object.py b/pyanalyze/type_object.py
index 0b1125d..7f88599 100644
--- a/pyanalyze/type_object.py
+++ b/pyanalyze/type_object.py
@@ -8,7 +8,18 @@ import inspect
 from typing import Container, Set, Sequence, Union
 from unittest import mock
 
-from .safe import safe_isinstance, safe_issubclass, is_typing_name
+from .safe import safe_isinstance, safe_issubclass, safe_in
+from .value import (
+    UNINITIALIZED_VALUE,
+    CanAssign,
+    CanAssignContext,
+    CanAssignError,
+    KnownValue,
+    TypedValue,
+    Value,
+    stringify_object,
+    unify_typevar_maps,
+)
 
 
 def get_mro(typ: Union[type, super]) -> Sequence[type]:
@@ -28,6 +39,7 @@ class TypeObject:
     typ: Union[type, super, str]
     base_classes: Set[Union[type, str]] = field(default_factory=set)
     is_protocol: bool = False
+    protocol_members: Set[str] = field(default_factory=set)
     is_thrift_enum: bool = field(init=False)
     is_universally_assignable: bool = field(init=False)
 
@@ -35,9 +47,6 @@ class TypeObject:
         if isinstance(self.typ, str):
             # Synthetic type
             self.is_universally_assignable = False
-            self.is_protocol = any(
-                is_typing_name(base, "Protocol") for base in self.base_classes
-            )
             self.is_thrift_enum = False
             return
         if isinstance(self.typ, super):
@@ -73,6 +82,56 @@ class TypeObject:
             )
         return self.is_assignable_to_type(other.typ)
 
+    def can_assign(
+        self,
+        self_val: Value,
+        other_val: Union[KnownValue, TypedValue],
+        ctx: CanAssignContext,
+    ) -> CanAssign:
+        other = other_val.get_type_object(ctx)
+        if other.is_universally_assignable:
+            return {}
+        if isinstance(self.typ, super):
+            if isinstance(other.typ, super):
+                return {}
+            return CanAssignError(f"Cannot assign to super object {self}")
+        if not self.is_protocol:
+            if other.is_protocol:
+                if self.typ is object:
+                    return {}
+                return CanAssignError(
+                    f"Cannot assign protocol {other} to non-protocol {self}"
+                )
+            if isinstance(self.typ, str):
+                if safe_in(self.typ, other.base_classes):
+                    return {}
+                return CanAssignError(f"Cannot assign {other} to {self}")
+            else:
+                for base in other.base_classes:
+                    if base is self.typ:
+                        return {}
+                    if isinstance(base, type) and safe_issubclass(base, self.typ):
+                        return {}
+                return CanAssignError(f"Cannot assign {other} to {self}")
+        else:
+            if isinstance(other.typ, super):
+                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)
+            return unify_typevar_maps(tv_maps)
+
     def is_instance(self, obj: object) -> bool:
         """Whether obj is an instance of this type."""
         return safe_isinstance(obj, self.typ)
@@ -86,3 +145,9 @@ class TypeObject:
         else:
             # TODO handle this for synthetic types (if necessary)
             return False
+
+    def __str__(self) -> str:
+        if isinstance(self.typ, str):
+            return self.typ
+        else:
+            return stringify_object(self.typ)
diff --git a/pyanalyze/typeshed.py b/pyanalyze/typeshed.py
index d93b5d8..0e5fb09 100644
--- a/pyanalyze/typeshed.py
+++ b/pyanalyze/typeshed.py
@@ -40,6 +40,7 @@ import inspect
 import sys
 from types import GeneratorType
 from typing import (
+    Set,
     Tuple,
     cast,
     Any,
@@ -254,6 +255,17 @@ class TypeshedFinder(object):
                     return True
         return False
 
+    def get_all_attributes(self, typ: Union[type, str]) -> Set[str]:
+        if isinstance(typ, str):
+            fq_name = typ
+        else:
+            fq_name = self._get_fq_name(typ)
+            if fq_name is None:
+                return set()
+        info = self._get_info_for_name(fq_name)
+        mod, _ = fq_name.rsplit(".", maxsplit=1)
+        return self._get_all_attributes_from_info(info, mod)
+
     def has_stubs(self, typ: type) -> bool:
         fq_name = self._get_fq_name(typ)
         if fq_name is None:
@@ -338,6 +350,29 @@ class TypeshedFinder(object):
         mod, _ = fq_name.rsplit(".", maxsplit=1)
         return self._has_attribute_from_info(info, mod, attr)
 
+    def _get_all_attributes_from_info(
+        self, info: typeshed_client.resolver.ResolvedName, mod: str
+    ) -> Set[str]:
+        if info is None:
+            return set()
+        elif isinstance(info, typeshed_client.ImportedInfo):
+            return self._get_all_attributes_from_info(
+                info.info, ".".join(info.source_module)
+            )
+        elif isinstance(info, typeshed_client.NameInfo):
+            if isinstance(info.ast, ast3.ClassDef):
+                if info.child_nodes is not None:
+                    return set(info.child_nodes)
+            elif isinstance(info.ast, ast3.Assign):
+                val = self._parse_expr(info.ast.value, mod)
+                if isinstance(val, KnownValue) and isinstance(val.val, type):
+                    return self.get_all_attributes(val.val)
+                else:
+                    return set()
+            else:
+                return set()
+        return set()
+
     def _has_attribute_from_info(
         self, info: typeshed_client.resolver.ResolvedName, mod: str, attr: str
     ) -> bool:
diff --git a/pyanalyze/value.py b/pyanalyze/value.py
index 2997b20..51edd71 100644
--- a/pyanalyze/value.py
+++ b/pyanalyze/value.py
@@ -46,7 +46,6 @@ import pyanalyze
 from pyanalyze.extensions import CustomCheck
 
 from .safe import safe_issubclass
-from .type_object import TypeObject
 
 T = TypeVar("T")
 # __builtin__ in Python 2 and builtins in Python 3
@@ -160,7 +159,9 @@ class CanAssignContext:
 
     """
 
-    def make_type_object(self, typ: Union[type, super, str]) -> TypeObject:
+    def make_type_object(
+        self, typ: Union[type, super, str]
+    ) -> "pyanalyze.type_object.TypeObject":
         """Return a :class:`pyanalyze.type_object.TypeObject` for this concrete type."""
         raise NotImplementedError
 
@@ -201,6 +202,9 @@ class CanAssignContext:
         """
         return None
 
+    def get_attribute_from_value(self, root_value: "Value", attribute: str) -> "Value":
+        return UNINITIALIZED_VALUE
+
 
 @dataclass(frozen=True)
 class CanAssignError:
@@ -347,10 +351,12 @@ class KnownValue(Value):
     def get_type(self) -> type:
         return type(self.val)
 
-    def get_type_object(self, ctx: Optional[CanAssignContext] = None) -> TypeObject:
+    def get_type_object(
+        self, ctx: Optional[CanAssignContext] = None
+    ) -> "pyanalyze.type_object.TypeObject":
         if ctx is not None:
             return ctx.make_type_object(type(self.val))
-        return TypeObject(type(self.val))
+        return pyanalyze.type_object.TypeObject(type(self.val))
 
     def get_type_value(self) -> Value:
         return KnownValue(type(self.val))
@@ -490,15 +496,17 @@ class TypedValue(Value):
 
     typ: Union[type, str]
     """The underlying type, or a fully qualified reference to one."""
-    _type_object: Optional[TypeObject] = field(
+    _type_object: Optional["pyanalyze.type_object.TypeObject"] = field(
         init=False, repr=False, hash=False, compare=False, default=None
     )
 
-    def get_type_object(self, ctx: Optional[CanAssignContext] = None) -> TypeObject:
+    def get_type_object(
+        self, ctx: Optional[CanAssignContext] = None
+    ) -> "pyanalyze.type_object.TypeObject":
         if self._type_object is None:
             if ctx is None:
                 # TODO: remoove this behavior and make ctx required
-                return TypeObject(self.typ)
+                return pyanalyze.type_object.TypeObject(self.typ)
             self._type_object = ctx.make_type_object(self.typ)
         return self._type_object
 
@@ -511,11 +519,9 @@ class TypedValue(Value):
         elif isinstance(other, KnownValue):
             if self_tobj.is_instance(other.val):
                 return {}
-            if other.get_type_object(ctx).is_assignable_to_type_object(self_tobj):
-                return {}
+            return self_tobj.can_assign(self, other, ctx)
         elif isinstance(other, TypedValue):
-            if other.get_type_object(ctx).is_assignable_to_type_object(self_tobj):
-                return {}
+            return self_tobj.can_assign(self, other, ctx)
         elif isinstance(other, SubclassValue):
             if isinstance(other.typ, TypedValue) and isinstance(
                 other.typ.typ, self.typ
@@ -541,10 +547,9 @@ class TypedValue(Value):
                 return {}
         elif isinstance(other, TypedValue):
             tobj = other.get_type_object(ctx)
-            if tobj.is_assignable_to_type_object(
-                self.get_type_object(ctx)
-            ) or tobj.is_assignable_to_type(int):
+            if tobj.is_assignable_to_type(int):
                 return {}
+            return self.get_type_object(ctx).can_assign(self, other, ctx)
         elif isinstance(other, MultiValuedValue):
             tv_maps = []
             for val in other.vals:
@@ -729,9 +734,8 @@ class SequenceIncompleteValue(GenericValue):
 
     def can_assign(self, other: Value, ctx: CanAssignContext) -> CanAssign:
         if isinstance(other, SequenceIncompleteValue):
-            if not other.get_type_object(ctx).is_assignable_to_type_object(
-                self.get_type_object(ctx)
-            ):
+            tv_map = self.get_type_object(ctx).can_assign(self, other, ctx)
+            if isinstance(tv_map, CanAssignError):
                 return CanAssignError(
                     f"Cannot assign {stringify_object(other.typ)} to"
                     f" {stringify_object(self.typ)}"
@@ -989,6 +993,8 @@ class CallableValue(TypedValue):
             signature = ctx.signature_from_value(other)
             if signature is None:
                 return CanAssignError(f"{other} is not a callable type")
+            if isinstance(signature, pyanalyze.signature.BoundMethodSignature):
+                signature = signature.get_signature()
             if isinstance(signature, pyanalyze.signature.Signature):
                 return self.signature.can_assign(signature, ctx)
 
@@ -1061,10 +1067,9 @@ class SubclassValue(Value):
             return self.typ.can_assign(other.typ, ctx)
         elif isinstance(other, KnownValue):
             if isinstance(other.val, type):
-                if isinstance(self.typ, TypedValue) and ctx.make_type_object(
-                    other.val
-                ).is_assignable_to_type_object(self.typ.get_type_object(ctx)):
-                    return {}
+                if isinstance(self.typ, TypedValue):
+                    self_tobj = self.typ.get_type_object(ctx)
+                    return self_tobj.can_assign(self, TypedValue(other.val), ctx)
                 elif isinstance(self.typ, TypeVarValue):
                     return {self.typ.typevar: TypedValue(other.val)}
         elif isinstance(other, TypedValue):


Test patch:
diff --git a/pyanalyze/test_type_object.py b/pyanalyze/test_type_object.py
index 8888951..b4bd689 100644
--- a/pyanalyze/test_type_object.py
+++ b/pyanalyze/test_type_object.py
@@ -45,9 +45,111 @@ class TestSyntheticType(TestNameCheckVisitorBase):
 
     @assert_passes()
     def test_protocol(self):
+        # Note that csv.writer expects this protocol:
+        # class _Writer(Protocol):
+        #    def write(self, s: str) -> Any: ...
         import csv
         import io
 
-        def capybara():
+        class BadWrite:
+            def write(self, s: int) -> object:
+                return object()
+
+        class GoodWrite:
+            def write(self, s: str) -> object:
+                return object()
+
+        class BadArgName:
+            def write(self, st: str) -> object:
+                return object()
+
+        def capybara(s: str):
             writer = io.StringIO()
             assert_is_value(csv.writer(writer), TypedValue("_csv._writer"))
+
+            csv.writer(1)  # E: incompatible_argument
+            csv.writer(s)  # E: incompatible_argument
+            csv.writer(BadWrite())  # E: incompatible_argument
+            csv.writer(GoodWrite())
+            csv.writer(BadArgName())  # E: incompatible_argument
+
+    @assert_passes()
+    def test_custom_subclasscheck(self):
+        class _ThriftEnumMeta(type):
+            def __subclasscheck__(self, subclass):
+                return hasattr(subclass, "_VALUES_TO_NAMES")
+
+        class ThriftEnum(metaclass=_ThriftEnumMeta):
+            pass
+
+        class IsOne:
+            _VALUES_TO_NAMES = {}
+
+        class IsntOne:
+            _NAMES_TO_VALUES = {}
+
+        def want_enum(te: ThriftEnum) -> None:
+            pass
+
+        def capybara(good_instance: IsOne, bad_instance: IsntOne, te: ThriftEnum):
+            want_enum(good_instance)
+            want_enum(bad_instance)  # E: incompatible_argument
+            want_enum(IsOne())
+            want_enum(IsntOne())  # E: incompatible_argument
+            want_enum(te)
+
+    @assert_passes()
+    def test_generic_stubonly(self):
+        import pkgutil
+
+        # pkgutil.read_code requires SupportsRead[bytes]
+
+        class Good:
+            def read(self, length: int = 0) -> bytes:
+                return b""
+
+        class Bad:
+            def read(self, length: int = 0) -> str:
+                return ""
+
+        def capybara():
+            pkgutil.read_code(1)  # E: incompatible_argument
+            pkgutil.read_code(Good())
+            pkgutil.read_code(Bad())  # E: incompatible_argument
+
+    @assert_passes()
+    def test_protocol_inheritance(self):
+        import cgi
+
+        # cgi.parse requires SupportsItemAccess[str, str]
+
+        class Good:
+            def __contains__(self, obj: object) -> bool:
+                return False
+
+            def __getitem__(self, k: str) -> str:
+                raise KeyError(k)
+
+            def __setitem__(self, k: str, v: str) -> None:
+                pass
+
+            def __delitem__(self, v: str) -> None:
+                pass
+
+        class Bad:
+            def __contains__(self, obj: object) -> bool:
+                return False
+
+            def __getitem__(self, k: bytes) -> str:
+                raise KeyError(k)
+
+            def __setitem__(self, k: str, v: str) -> None:
+                pass
+
+            def __delitem__(self, v: str) -> None:
+                pass
+
+        def capybara():
+            cgi.parse(environ=Good())
+            cgi.parse(environ=Bad())  # E: incompatible_argument
+            cgi.parse(environ=1)  # E: incompatible_argument