Skip to content

Commit 477638f

Browse files
authored
[ty] Compare protocol receiver restrictions before binding (astral-sh#29443)
Improve compatibility checks for protocol methods with explicit receiver restrictions. When both receivers specialize the same protocol, we can use the variance of the protocol's typevars to prove that the target restriction implies the implementation's restriction. This lets us avoid recursive receiver checks that introduce fresh type variables. We preserve the existing structural fallback when the variance-based proof cannot establish compatibility. Pulled this out of astral-sh#29027; it claws back one of the larger perf regressions from that PR, but is a standalone fix that can be reviewed/merged in isolation.
1 parent 9cced9a commit 477638f

5 files changed

Lines changed: 360 additions & 5 deletions

File tree

‎crates/ty_python_semantic/resources/mdtest/generics/legacy/classes.md‎

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3727,6 +3727,122 @@ exact_class: type[Just[object]] = object # no diagnostic
37273727
subclass_class: type[Just[object]] = Custom # error: [invalid-assignment]
37283728
```
37293729

3730+
## Recursive protocol receiver restrictions
3731+
3732+
An implementation can use the same receiver restriction as a generic protocol. Its other parameters
3733+
and return types must still match.
3734+
3735+
```py
3736+
from __future__ import annotations
3737+
from collections.abc import Callable, Iterable
3738+
from typing import Generic, Protocol, TypeVar
3739+
3740+
T = TypeVar("T")
3741+
S = TypeVar("S")
3742+
3743+
class Chain(Protocol[T]):
3744+
def value(self) -> T:
3745+
raise RuntimeError
3746+
def step(self: Chain[tuple[S]], callback: Callable[[S], T]) -> None:
3747+
raise RuntimeError
3748+
3749+
class Inherited(Chain[T]):
3750+
def __init__(self, values: Iterable[T]) -> None: ...
3751+
3752+
class Structural(Generic[T]):
3753+
def value(self) -> T:
3754+
raise NotImplementedError
3755+
def step(self: Chain[tuple[S]], callback: Callable[[S], T]) -> None:
3756+
raise NotImplementedError
3757+
3758+
class Invalid(Generic[T]):
3759+
def value(self) -> T:
3760+
raise NotImplementedError
3761+
def step(self: Chain[tuple[S]], callback: Callable[[S], str]) -> None:
3762+
raise NotImplementedError
3763+
3764+
def accepts(value: Chain[int]) -> None: ...
3765+
def check(inherited: Inherited[int], structural: Structural[int], wrong: Structural[str], invalid: Invalid[int]):
3766+
accepts(inherited) # no diagnostic
3767+
accepts(structural) # no diagnostic
3768+
accepts(wrong) # error: [invalid-argument-type]
3769+
accepts(invalid) # error: [invalid-argument-type]
3770+
```
3771+
3772+
A subprotocol must pass the same assignability check.
3773+
3774+
```py
3775+
class StructuralProtocol(Protocol[T]):
3776+
def value(self) -> T: ...
3777+
def step(self: Chain[tuple[S]], callback: Callable[[S], T]) -> None: ...
3778+
3779+
def check_protocol(valid: StructuralProtocol[int], invalid: StructuralProtocol[str]):
3780+
accepts(valid) # no diagnostic
3781+
accepts(invalid) # error: [invalid-argument-type]
3782+
```
3783+
3784+
The expected protocol also supplies the inherited constructor's type argument.
3785+
3786+
```py
3787+
contextual: Chain[int] = Inherited(()) # no diagnostic
3788+
reveal_type(contextual) # revealed: Inherited[int]
3789+
```
3790+
3791+
## Variance of protocol receiver restrictions
3792+
3793+
An implementation can accept a broader receiver than a protocol requires.
3794+
3795+
```py
3796+
from __future__ import annotations
3797+
from typing import Protocol, TypeVar
3798+
3799+
T_co = TypeVar("T_co", covariant=True)
3800+
3801+
class Reader(Protocol[T_co]):
3802+
def read(self) -> T_co: ...
3803+
3804+
class Required(Protocol):
3805+
def method(self: Reader[bool]) -> None: ...
3806+
3807+
class Broad:
3808+
def method(self: Reader[int]) -> None: ...
3809+
3810+
class Incompatible:
3811+
def method(self: Reader[str]) -> None: ...
3812+
3813+
def accepts(value: Required) -> None: ...
3814+
def check(broad: Broad, incompatible: Incompatible):
3815+
accepts(broad) # no diagnostic
3816+
accepts(incompatible) # error: [invalid-argument-type]
3817+
```
3818+
3819+
## Structurally equivalent protocol receiver restrictions
3820+
3821+
Receiver annotations can match structurally even when their type arguments do not satisfy the
3822+
protocol's declared variance.
3823+
3824+
```py
3825+
from __future__ import annotations
3826+
from typing import Protocol, TypeVar
3827+
3828+
T_co = TypeVar("T_co", covariant=True)
3829+
3830+
class Recursive(Protocol[T_co]):
3831+
def child(self) -> Recursive[T_co]: ...
3832+
3833+
class Required(Protocol):
3834+
def method(self: Recursive[int]) -> None: ...
3835+
3836+
class Implementation:
3837+
def child(self) -> Implementation:
3838+
return self
3839+
def method(self: Recursive[str]) -> None: ...
3840+
3841+
def accepts(value: Required) -> None: ...
3842+
def check(value: Implementation):
3843+
accepts(value) # no diagnostic
3844+
```
3845+
37303846
## Aliased `Self` in explicit receivers
37313847

37323848
Specializing a generic class also specializes the upper bound of `Self` inside type alias arguments.

‎crates/ty_python_semantic/resources/mdtest/generics/pep695/classes.md‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2834,6 +2834,90 @@ def invalid() -> Nested:
28342834
return Box("wrong") # error: [invalid-return-type]
28352835
```
28362836

2837+
## Recursive protocol receiver restrictions
2838+
2839+
An implementation can use the same receiver restriction as a generic protocol. Its other parameters
2840+
and return types must still match.
2841+
2842+
```py
2843+
from __future__ import annotations
2844+
from collections.abc import Callable, Iterable
2845+
from typing import Protocol
2846+
2847+
class Chain[T](Protocol):
2848+
def value(self) -> T:
2849+
raise RuntimeError
2850+
def step[S](self: Chain[tuple[S]], callback: Callable[[S], T]) -> None:
2851+
raise RuntimeError
2852+
2853+
class Inherited[T](Chain[T]):
2854+
def __init__(self, values: Iterable[T]) -> None: ...
2855+
2856+
class Structural[T]:
2857+
def value(self) -> T:
2858+
raise NotImplementedError
2859+
def step[S](self: Chain[tuple[S]], callback: Callable[[S], T]) -> None:
2860+
raise NotImplementedError
2861+
2862+
class Invalid[T]:
2863+
def value(self) -> T:
2864+
raise NotImplementedError
2865+
def step[S](self: Chain[tuple[S]], callback: Callable[[S], str]) -> None:
2866+
raise NotImplementedError
2867+
2868+
def accepts(value: Chain[int]) -> None: ...
2869+
def check(inherited: Inherited[int], structural: Structural[int], wrong: Structural[str], invalid: Invalid[int]):
2870+
accepts(inherited) # no diagnostic
2871+
accepts(structural) # no diagnostic
2872+
accepts(wrong) # error: [invalid-argument-type]
2873+
accepts(invalid) # error: [invalid-argument-type]
2874+
```
2875+
2876+
A subprotocol must pass the same assignability check.
2877+
2878+
```py
2879+
class StructuralProtocol[T](Protocol):
2880+
def value(self) -> T: ...
2881+
def step[S](self: Chain[tuple[S]], callback: Callable[[S], T]) -> None: ...
2882+
2883+
def check_protocol(valid: StructuralProtocol[int], invalid: StructuralProtocol[str]):
2884+
accepts(valid) # no diagnostic
2885+
accepts(invalid) # error: [invalid-argument-type]
2886+
```
2887+
2888+
The expected protocol also supplies the inherited constructor's type argument.
2889+
2890+
```py
2891+
contextual: Chain[int] = Inherited(()) # no diagnostic
2892+
reveal_type(contextual) # revealed: Inherited[int]
2893+
```
2894+
2895+
## Variance of protocol receiver restrictions
2896+
2897+
An implementation can accept a broader receiver than a protocol requires.
2898+
2899+
```py
2900+
from __future__ import annotations
2901+
from typing import Protocol
2902+
2903+
class Reader[T](Protocol):
2904+
def read(self) -> T: ...
2905+
2906+
class Required(Protocol):
2907+
def method(self: Reader[bool]) -> None: ...
2908+
2909+
class Broad:
2910+
def method(self: Reader[int]) -> None: ...
2911+
2912+
class Incompatible:
2913+
def method(self: Reader[str]) -> None: ...
2914+
2915+
def accepts(value: Required) -> None: ...
2916+
def check(broad: Broad, incompatible: Incompatible):
2917+
accepts(broad) # no diagnostic
2918+
accepts(incompatible) # error: [invalid-argument-type]
2919+
```
2920+
28372921
## Aliased `Self` in explicit receivers
28382922

28392923
Specializing a generic class also specializes the upper bound of `Self` inside type alias arguments.

‎crates/ty_python_semantic/resources/mdtest/protocols.md‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4430,6 +4430,27 @@ static_assert(is_subtype_of(Text, ConsoleRenderable))
44304430
static_assert(is_assignable_to(Text, ConsoleRenderable))
44314431
```
44324432

4433+
## Receiver parameter names
4434+
4435+
Receiver parameter names do not affect method compatibility, including methods with explicit
4436+
protocol receiver restrictions.
4437+
4438+
```py
4439+
from typing import Protocol
4440+
4441+
class Receiver(Protocol):
4442+
value: int
4443+
4444+
class Required(Protocol):
4445+
def method(self: Receiver) -> None: ...
4446+
4447+
class Implementation:
4448+
def method(receiver: Receiver) -> None: ...
4449+
4450+
def check(value: Implementation) -> Required:
4451+
return value # no diagnostic
4452+
```
4453+
44334454
## Recursive protocol receiver binding
44344455

44354456
A classmethod on a generic protocol can cause receiver binding for another method to depend on

‎crates/ty_python_semantic/src/types/protocol_class.rs‎

Lines changed: 90 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ use crate::types::attribute_write::{
1212
property_setter_value_type,
1313
};
1414
use crate::types::overrides::{VariableKind, effective_superclass_variable_kind};
15-
use crate::types::relation::{DisjointnessChecker, TypeRelationChecker};
15+
use crate::types::relation::{DisjointnessChecker, TypeRelationChecker, TypeVarEvaluation};
1616
use crate::types::visitor::any_over_type_expanding_aliases;
1717
use crate::types::{TypeContext, UpcastPolicy};
1818
use crate::{
@@ -1451,6 +1451,19 @@ struct ProtocolMemberAccess<'a, 'db> {
14511451
}
14521452

14531453
impl<'a, 'db> ProtocolMemberAccess<'a, 'db> {
1454+
fn unbound_instance_method(self, db: &'db dyn Db) -> Option<&'db CallableSignature<'db>> {
1455+
// Raw declaration signatures do not include a pending materialization.
1456+
if self.mode != ProtocolMemberAccessMode::Instance || self.materialization.is_some() {
1457+
return None;
1458+
}
1459+
let ProtocolMemberKind::Method(Type::Callable(callable), ProtocolMethodKind::Instance) =
1460+
self.declaration.kind
1461+
else {
1462+
return None;
1463+
};
1464+
Some(callable.signatures(db))
1465+
}
1466+
14541467
fn read(self) -> Option<ProtocolMemberReadAccess<'a, 'db>> {
14551468
let supported = match self.declaration.kind {
14561469
ProtocolMemberKind::Method(..) => true,
@@ -2596,6 +2609,46 @@ fn protocol_member_read_type<'db>(
25962609
}
25972610

25982611
impl<'c, 'db> TypeRelationChecker<'_, 'c, 'db> {
2612+
/// Prove receiver compatibility without unfolding two instances of the same protocol.
2613+
///
2614+
/// A target receiver restriction `P[B]` implies a source restriction `P[A]` if `P[B]`
2615+
/// is assignable to `P[A]`. Comparing their type arguments using the protocol's variance
2616+
/// is sufficient; checking the implementing instance against both annotations would
2617+
/// recursively check this same method with fresh type variables. Retaining the receiver
2618+
/// parameter keeps its constraints in the same scope as the other parameters and return type.
2619+
fn try_check_protocol_receiver_pair(
2620+
&self,
2621+
db: &'db dyn Db,
2622+
source: &CallableSignature<'db>,
2623+
target: &CallableSignature<'db>,
2624+
source_self_type: Type<'db>,
2625+
target_self_type: Type<'db>,
2626+
) -> Option<ConstraintSet<'db, 'c>> {
2627+
let target_receiver = target.overloads.first()?.parameters().get_positional(0)?;
2628+
let receiver_protocol = target_receiver
2629+
.annotated_type()
2630+
.resolve_type_alias(db)
2631+
.as_protocol_instance()?;
2632+
let origin = receiver_protocol.class_origin(db)?.class_literal(db);
2633+
let prepare = |signatures: &CallableSignature<'db>, self_type| {
2634+
signatures
2635+
.iter()
2636+
.map(|signature| {
2637+
signature.with_nominal_protocol_receiver(db, self.env, origin, self_type)
2638+
})
2639+
.collect::<Option<Vec<_>>>()
2640+
.map(CallableSignature::from_overloads)
2641+
};
2642+
let source = prepare(source, source_self_type)?;
2643+
let target = prepare(target, target_self_type)?;
2644+
let mut checker = self.clone();
2645+
checker.typevar_evaluation = TypeVarEvaluation::Lazy;
2646+
let result = checker.check_callable_signature_pair(db, &source, &target);
2647+
// This is a sufficient proof. A failed nominal comparison does not rule out a
2648+
// structural match, even between specializations of the same protocol.
2649+
(!result.is_never_satisfied(db, self.env, self.inferable)).then_some(result)
2650+
}
2651+
25992652
fn check_protocol_member_read(
26002653
&self,
26012654
db: &'db dyn Db,
@@ -2667,6 +2720,36 @@ impl<'c, 'db> TypeRelationChecker<'_, 'c, 'db> {
26672720
(implementation_self_binding_ty, protocol_self_binding_ty)
26682721
};
26692722

2723+
if let Some(target) = required.unbound_instance_method(db)
2724+
&& let Type::BoundMethod(method) = attribute_type
2725+
&& method.signature_receiver(db) == protocol_receiver_binding_ty
2726+
&& let Some(source) = method.unbound_signatures(db)
2727+
{
2728+
// A stored bound method keeps its captured receiver; only descriptor binding to
2729+
// this object lets us use the target's receiver restriction as a premise.
2730+
let source_descriptor = ty
2731+
.class_member_with_policy(
2732+
db,
2733+
env,
2734+
member.name,
2735+
MemberLookupPolicy::NO_INSTANCE_FALLBACK,
2736+
)
2737+
.place
2738+
.ignore_possibly_undefined();
2739+
if source_descriptor
2740+
.is_some_and(|descriptor| descriptor.function_like_kind(db).is_some())
2741+
&& let Some(result) = self.try_check_protocol_receiver_pair(
2742+
db,
2743+
source,
2744+
target,
2745+
method.typing_self_type(db),
2746+
protocol_self_binding_ty,
2747+
)
2748+
{
2749+
return result;
2750+
}
2751+
}
2752+
26702753
let Some(Type::Callable(required_callable)) = required
26712754
.read()
26722755
.and_then(|read| read.result_type(db, env, None))
@@ -3006,6 +3089,12 @@ impl<'c, 'db> TypeRelationChecker<'_, 'c, 'db> {
30063089
self.always()
30073090
} else if source.read().is_none() {
30083091
self.never()
3092+
} else if let Some(source) = source.unbound_instance_method(db)
3093+
&& let Some(target) = target.unbound_instance_method(db)
3094+
&& let Some(result) =
3095+
self.try_check_protocol_receiver_pair(db, source, target, source_type, source_type)
3096+
{
3097+
result
30093098
} else {
30103099
let bind_read = |access: ProtocolMemberAccess<'_, 'db>,
30113100
member: &ProtocolMember<'_, 'db>| {

0 commit comments

Comments
 (0)