Infer specialization of protocol parameter

Author: dcreager

crates/ty_python_semantic/resources/mdtest/generics/legacy/functions.md

@@ -66,21 +66,42 @@ reveal_type(f("string"))  # revealed: Literal["string"]
 ## Inferring “deep” generic parameter types
 
 The matching up of call arguments and discovery of constraints on typevars can be a recursive
-process for arbitrarily-nested generic types in parameters.
+process for arbitrarily-nested generic classes and protocols in parameters.
+
+TODO: Note that we can currently only infer a specialization for a generic protocol when the
+argument _explicitly_ implements the protocol by listing it as a base class.
 
 ```py
-from typing import TypeVar
+from typing import Protocol, TypeVar
 
 T = TypeVar("T")
 
-def f(x: list[T]) -> T:
+class CanIndex(Protocol[T]):
+    def __getitem__(self, index: int) -> T: ...
+
+class ExplicitlyImplements(CanIndex[T]): ...
+
+def takes_in_list(x: list[T]) -> T:
+    return x[0]
+
+def takes_in_protocol(x: CanIndex[T]) -> T:
     return x[0]
 
-def deep(x: list[str]) -> None:
-    reveal_type(f(x))  # revealed: str
+def deep_list(x: list[str]) -> None:
+    reveal_type(takes_in_list(x))  # revealed: str
+    # TODO: revealed: str
+    reveal_type(takes_in_protocol(x))  # revealed: Unknown
 
-def deeper(x: list[set[str]]) -> None:
-    reveal_type(f(x))  # revealed: set[str]
+def deeper_list(x: list[set[str]]) -> None:
+    reveal_type(takes_in_list(x))  # revealed: set[str]
+    # TODO: revealed: set[str]
+    reveal_type(takes_in_protocol(x))  # revealed: Unknown
+
+def deep_explicit(x: ExplicitlyImplements[str]) -> None:
+    reveal_type(takes_in_protocol(x))  # revealed: str
+
+def deeper_explicit(x: ExplicitlyImplements[set[str]]) -> None:
+    reveal_type(takes_in_protocol(x))  # revealed: set[str]
 ```
 
 This also works when passing in arguments that are subclasses of the parameter type.
@@ -89,8 +110,19 @@ This also works when passing in arguments that are subclasses of the parameter t
 class Sub(list[int]): ...
 class GenericSub(list[T]): ...
 
-reveal_type(f(Sub()))  # revealed: int
-reveal_type(f(GenericSub[str]()))  # revealed: str
+reveal_type(takes_in_list(Sub()))  # revealed: int
+# TODO: revealed: int
+reveal_type(takes_in_protocol(Sub()))  # revealed: Unknown
+
+reveal_type(takes_in_list(GenericSub[str]()))  # revealed: str
+# TODO: revealed: str
+reveal_type(takes_in_protocol(GenericSub[str]()))  # revealed: Unknown
+
+class ExplicitSub(ExplicitlyImplements[int]): ...
+class ExplicitGenericSub(ExplicitlyImplements[T]): ...
+
+reveal_type(takes_in_protocol(ExplicitSub()))  # revealed: int
+reveal_type(takes_in_protocol(ExplicitGenericSub[str]()))  # revealed: str
 ```
 
 ## Inferring a bound typevar

crates/ty_python_semantic/resources/mdtest/generics/pep695/functions.md

@@ -61,17 +61,42 @@ reveal_type(f("string"))  # revealed: Literal["string"]
 ## Inferring “deep” generic parameter types
 
 The matching up of call arguments and discovery of constraints on typevars can be a recursive
-process for arbitrarily-nested generic types in parameters.
+process for arbitrarily-nested generic classes and protocols in parameters.
+
+TODO: Note that we can currently only infer a specialization for a generic protocol when the
+argument _explicitly_ implements the protocol by listing it as a base class.
 
 ```py
-def f[T](x: list[T]) -> T:
+from typing import Protocol, TypeVar
+
+S = TypeVar("S")
+
+class CanIndex(Protocol[S]):
+    def __getitem__(self, index: int) -> S: ...
+
+class ExplicitlyImplements[T](CanIndex[T]): ...
+
+def takes_in_list[T](x: list[T]) -> T:
+    return x[0]
+
+def takes_in_protocol[T](x: CanIndex[T]) -> T:
     return x[0]
 
-def deep(x: list[str]) -> None:
-    reveal_type(f(x))  # revealed: str
+def deep_list(x: list[str]) -> None:
+    reveal_type(takes_in_list(x))  # revealed: str
+    # TODO: revealed: str
+    reveal_type(takes_in_protocol(x))  # revealed: Unknown
+
+def deeper_list(x: list[set[str]]) -> None:
+    reveal_type(takes_in_list(x))  # revealed: set[str]
+    # TODO: revealed: set[str]
+    reveal_type(takes_in_protocol(x))  # revealed: Unknown
+
+def deep_explicit(x: ExplicitlyImplements[str]) -> None:
+    reveal_type(takes_in_protocol(x))  # revealed: str
 
-def deeper(x: list[set[str]]) -> None:
-    reveal_type(f(x))  # revealed: set[str]
+def deeper_explicit(x: ExplicitlyImplements[set[str]]) -> None:
+    reveal_type(takes_in_protocol(x))  # revealed: set[str]
 ```
 
 This also works when passing in arguments that are subclasses of the parameter type.
@@ -80,8 +105,19 @@ This also works when passing in arguments that are subclasses of the parameter t
 class Sub(list[int]): ...
 class GenericSub[T](list[T]): ...
 
-reveal_type(f(Sub()))  # revealed: int
-reveal_type(f(GenericSub[str]()))  # revealed: str
+reveal_type(takes_in_list(Sub()))  # revealed: int
+# TODO: revealed: int
+reveal_type(takes_in_protocol(Sub()))  # revealed: Unknown
+
+reveal_type(takes_in_list(GenericSub[str]()))  # revealed: str
+# TODO: revealed: str
+reveal_type(takes_in_protocol(GenericSub[str]()))  # revealed: Unknown
+
+class ExplicitSub(ExplicitlyImplements[int]): ...
+class ExplicitGenericSub[T](ExplicitlyImplements[T]): ...
+
+reveal_type(takes_in_protocol(ExplicitSub()))  # revealed: int
+reveal_type(takes_in_protocol(ExplicitGenericSub[str]()))  # revealed: str
 ```
 
 ## Inferring a bound typevar

crates/ty_python_semantic/src/types/generics.rs

@@ -4,7 +4,7 @@ use rustc_hash::FxHashMap;
 use crate::semantic_index::SemanticIndex;
 use crate::types::class::ClassType;
 use crate::types::class_base::ClassBase;
-use crate::types::instance::NominalInstanceType;
+use crate::types::instance::{NominalInstanceType, Protocol, ProtocolInstanceType};
 use crate::types::signatures::{Parameter, Parameters, Signature};
 use crate::types::{
     declaration_type, todo_type, KnownInstanceType, Type, TypeVarBoundOrConstraints,
@@ -630,7 +630,10 @@ impl<'db> SpecializationBuilder<'db> {
         // ```
         //
         // without specializing `T` to `None`.
-        if !actual.is_never() && actual.is_subtype_of(self.db, formal) {
+        if !matches!(formal, Type::ProtocolInstance(_))
+            && !actual.is_never()
+            && actual.is_subtype_of(self.db, formal)
+        {
             return Ok(());
         }
 
@@ -678,6 +681,14 @@ impl<'db> SpecializationBuilder<'db> {
                 Type::NominalInstance(NominalInstanceType {
                     class: ClassType::Generic(formal_alias),
                     ..
+                })
+                // TODO: This will only handle classes that explicit implement a generic protocol
+                // by listing it as a base class. To handle classes that implicitly implement a
+                // generic protocol, we will need to check the types of the protocol members to be
+                // able to infer the specialization of the protocol that the class implements.
+                | Type::ProtocolInstance(ProtocolInstanceType {
+                    inner: Protocol::FromClass(ClassType::Generic(formal_alias)),
+                    ..
                 }),
                 Type::NominalInstance(NominalInstanceType {
                     class: actual_class,