Skip to content

Commit 25cbf38

Browse files
sharkdpsharkdp
authored
[ty] Patch Self for fallback-methods on NamedTuples and TypedDicts (#20328)
## Summary We use classes like [`_typeshed._type_checker_internals.NamedTupleFallback`](https://github.com/python/typeshed/blob/d9c76e1d9f0c0000c1971c3e936a69dd55f49ce1/stdlib/_typeshed/_type_checker_internals.pyi#L54-L75) to tack on additional attributes/methods to instances of user-defined `NamedTuple`s (or `TypedDict`s), even though these classes are not present in the MRO of those types. The problem is that those classes use implicit and explicit `Self` annotations which refer to `NamedTupleFallback` itself, instead of to the actual type that we're adding those methods to: ```py class NamedTupleFallback(tuple[Any, ...]): # […] def _replace(self, **kwargs: Any) -> typing_extensions.Self: ... ``` In effect, when we access `_replace` on an instance of a custom `NamedTuple` instance, its `self` parameter and return type refer to the wrong `Self`. This leads to incorrect *"Argument to bound method `_replace` is incorrect: Argument type `Person` does not satisfy upper bound `NamedTupleFallback` of type variable `Self`"* errors on #18007. It would also lead to similar errors on `TypedDict`s, if they would already implement assignability properly. ## Test Plan I applied the following patch to typeshed and verified that no errors appear anymore. <details> ```diff diff --git a/crates/ty_vendored/vendor/typeshed/stdlib/_typeshed/_type_checker_internals.pyi b/crates/ty_vendored/vendor/typeshed/stdlib/_typeshed/_type_checker_internals.pyi index feb22aa..8e41034f19 100644 --- a/crates/ty_vendored/vendor/typeshed/stdlib/_typeshed/_type_checker_internals.pyi +++ b/crates/ty_vendored/vendor/typeshed/stdlib/_typeshed/_type_checker_internals.pyi @@ -29,27 +29,27 @@ class TypedDictFallback(Mapping[str, object], metaclass=ABCMeta): __readonly_keys__: ClassVar[frozenset[str]] __mutable_keys__: ClassVar[frozenset[str]] - def copy(self) -> typing_extensions.Self: ... + def copy(self: typing_extensions.Self) -> typing_extensions.Self: ... # Using Never so that only calls using mypy plugin hook that specialize the signature # can go through. - def setdefault(self, k: Never, default: object) -> object: ... + def setdefault(self: typing_extensions.Self, k: Never, default: object) -> object: ... # Mypy plugin hook for 'pop' expects that 'default' has a type variable type. - def pop(self, k: Never, default: _T = ...) -> object: ... # pyright: ignore[reportInvalidTypeVarUse] - def update(self, m: typing_extensions.Self, /) -> None: ... - def __delitem__(self, k: Never) -> None: ... - def items(self) -> dict_items[str, object]: ... - def keys(self) -> dict_keys[str, object]: ... - def values(self) -> dict_values[str, object]: ... + def pop(self: typing_extensions.Self, k: Never, default: _T = ...) -> object: ... # pyright: ignore[reportInvalidTypeVarUse] + def update(self: typing_extensions.Self, m: typing_extensions.Self, /) -> None: ... + def __delitem__(self: typing_extensions.Self, k: Never) -> None: ... + def items(self: typing_extensions.Self) -> dict_items[str, object]: ... + def keys(self: typing_extensions.Self) -> dict_keys[str, object]: ... + def values(self: typing_extensions.Self) -> dict_values[str, object]: ... @overload - def __or__(self, value: typing_extensions.Self, /) -> typing_extensions.Self: ... + def __or__(self: typing_extensions.Self, value: typing_extensions.Self, /) -> typing_extensions.Self: ... @overload - def __or__(self, value: dict[str, Any], /) -> dict[str, object]: ... + def __or__(self: typing_extensions.Self, value: dict[str, Any], /) -> dict[str, object]: ... @overload - def __ror__(self, value: typing_extensions.Self, /) -> typing_extensions.Self: ... + def __ror__(self: typing_extensions.Self, value: typing_extensions.Self, /) -> typing_extensions.Self: ... @overload - def __ror__(self, value: dict[str, Any], /) -> dict[str, object]: ... + def __ror__(self: typing_extensions.Self, value: dict[str, Any], /) -> dict[str, object]: ... # supposedly incompatible definitions of __or__ and __ior__ - def __ior__(self, value: typing_extensions.Self, /) -> typing_extensions.Self: ... # type: ignore[misc] + def __ior__(self: typing_extensions.Self, value: typing_extensions.Self, /) -> typing_extensions.Self: ... # type: ignore[misc] # Fallback type providing methods and attributes that appear on all `NamedTuple` types. class NamedTupleFallback(tuple[Any, ...]): @@ -61,18 +61,18 @@ class NamedTupleFallback(tuple[Any, ...]): __orig_bases__: ClassVar[tuple[Any, ...]] @overload - def __init__(self, typename: str, fields: Iterable[tuple[str, Any]], /) -> None: ... + def __init__(self: typing_extensions.Self, typename: str, fields: Iterable[tuple[str, Any]], /) -> None: ... @overload @typing_extensions.deprecated( "Creating a typing.NamedTuple using keyword arguments is deprecated and support will be removed in Python 3.15" ) - def __init__(self, typename: str, fields: None = None, /, **kwargs: Any) -> None: ... + def __init__(self: typing_extensions.Self, typename: str, fields: None = None, /, **kwargs: Any) -> None: ... @classmethod def _make(cls, iterable: Iterable[Any]) -> typing_extensions.Self: ... - def _asdict(self) -> dict[str, Any]: ... - def _replace(self, **kwargs: Any) -> typing_extensions.Self: ... + def _asdict(self: typing_extensions.Self) -> dict[str, Any]: ... + def _replace(self: typing_extensions.Self, **kwargs: Any) -> typing_extensions.Self: ... if sys.version_info >= (3, 13): - def __replace__(self, **kwargs: Any) -> typing_extensions.Self: ... + def __replace__(self: typing_extensions.Self, **kwargs: Any) -> typing_extensions.Self: ... # Non-default variations to accommodate couroutines, and `AwaitableGenerator` having a 4th type parameter. _S = TypeVar("_S") ``` </details>
1 parent 9a9ebc3 commit 25cbf38

File tree

6 files changed

+188
-7
lines changed

6 files changed

+188
-7
lines changed

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

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -272,16 +272,34 @@ reveal_type(Person._make) # revealed: bound method <class 'Person'>._make(itera
272272
reveal_type(Person._asdict) # revealed: def _asdict(self) -> dict[str, Any]
273273
reveal_type(Person._replace) # revealed: def _replace(self, **kwargs: Any) -> Self@_replace
274274

275-
# TODO: should be `Person` once we support `Self`
275+
# TODO: should be `Person` once we support implicit type of `self`
276276
reveal_type(Person._make(("Alice", 42))) # revealed: Unknown
277277

278278
person = Person("Alice", 42)
279279

280280
reveal_type(person._asdict()) # revealed: dict[str, Any]
281-
# TODO: should be `Person` once we support `Self`
281+
# TODO: should be `Person` once we support implicit type of `self`
282282
reveal_type(person._replace(name="Bob")) # revealed: Unknown
283283
```
284284

285+
When accessing them on child classes of generic `NamedTuple`s, the return type is specialized
286+
accordingly:
287+
288+
```py
289+
from typing import NamedTuple, Generic, TypeVar
290+
291+
T = TypeVar("T")
292+
293+
class Box(NamedTuple, Generic[T]):
294+
content: T
295+
296+
class IntBox(Box[int]):
297+
pass
298+
299+
# TODO: should be `IntBox` once we support the implicit type of `self`
300+
reveal_type(IntBox(1)._replace(content=42)) # revealed: Unknown
301+
```
302+
285303
## `collections.namedtuple`
286304

287305
```py

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -627,6 +627,18 @@ alice: Employee = {"name": "Alice", "employee_id": 1}
627627

628628
# error: [missing-typed-dict-key] "Missing required key 'employee_id' in TypedDict `Employee` constructor"
629629
eve: Employee = {"name": "Eve"}
630+
631+
def combine(p: Person, e: Employee):
632+
# TODO: Should be `Person` once we support the implicit type of self
633+
reveal_type(p.copy()) # revealed: Unknown
634+
# TODO: Should be `Employee` once we support the implicit type of self
635+
reveal_type(e.copy()) # revealed: Unknown
636+
637+
reveal_type(p | p) # revealed: Person
638+
reveal_type(e | e) # revealed: Employee
639+
640+
# TODO: Should be `Person` once we support the implicit type of self and subtyping for TypedDicts
641+
reveal_type(p | e) # revealed: Employee
630642
```
631643

632644
When inheriting from a `TypedDict` with a different `total` setting, inherited fields maintain their

‎crates/ty_python_semantic/src/types.rs‎

Lines changed: 99 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5971,6 +5971,19 @@ impl<'db> Type<'db> {
59715971
self
59725972
}
59735973
}
5974+
TypeMapping::ReplaceSelf { new_upper_bound } => {
5975+
if bound_typevar.typevar(db).is_self(db) {
5976+
Type::TypeVar(
5977+
BoundTypeVarInstance::synthetic_self(
5978+
db,
5979+
*new_upper_bound,
5980+
bound_typevar.binding_context(db)
5981+
)
5982+
)
5983+
} else {
5984+
self
5985+
}
5986+
}
59745987
TypeMapping::PromoteLiterals | TypeMapping::BindLegacyTypevars(_) |
59755988
TypeMapping::MarkTypeVarsInferable(_) => self,
59765989
TypeMapping::Materialize(materialization_kind) => {
@@ -5994,7 +6007,8 @@ impl<'db> Type<'db> {
59946007
}
59956008
TypeMapping::PromoteLiterals |
59966009
TypeMapping::BindLegacyTypevars(_) |
5997-
TypeMapping::BindSelf(_)
6010+
TypeMapping::BindSelf(_) |
6011+
TypeMapping::ReplaceSelf { .. }
59986012
=> self,
59996013
TypeMapping::Materialize(materialization_kind) => Type::NonInferableTypeVar(bound_typevar.materialize_impl(db, *materialization_kind, visitor))
60006014

@@ -6008,6 +6022,7 @@ impl<'db> Type<'db> {
60086022
TypeMapping::PartialSpecialization(_) |
60096023
TypeMapping::PromoteLiterals |
60106024
TypeMapping::BindSelf(_) |
6025+
TypeMapping::ReplaceSelf { .. } |
60116026
TypeMapping::MarkTypeVarsInferable(_) |
60126027
TypeMapping::Materialize(_) => self,
60136028
}
@@ -6116,6 +6131,7 @@ impl<'db> Type<'db> {
61166131
TypeMapping::PartialSpecialization(_) |
61176132
TypeMapping::BindLegacyTypevars(_) |
61186133
TypeMapping::BindSelf(_) |
6134+
TypeMapping::ReplaceSelf { .. } |
61196135
TypeMapping::MarkTypeVarsInferable(_) |
61206136
TypeMapping::Materialize(_) => self,
61216137
TypeMapping::PromoteLiterals => self.literal_fallback_instance(db)
@@ -6127,6 +6143,7 @@ impl<'db> Type<'db> {
61276143
TypeMapping::PartialSpecialization(_) |
61286144
TypeMapping::BindLegacyTypevars(_) |
61296145
TypeMapping::BindSelf(_) |
6146+
TypeMapping::ReplaceSelf { .. } |
61306147
TypeMapping::MarkTypeVarsInferable(_) |
61316148
TypeMapping::PromoteLiterals => self,
61326149
TypeMapping::Materialize(materialization_kind) => match materialization_kind {
@@ -6662,6 +6679,8 @@ pub enum TypeMapping<'a, 'db> {
66626679
BindLegacyTypevars(BindingContext<'db>),
66636680
/// Binds any `typing.Self` typevar with a particular `self` class.
66646681
BindSelf(Type<'db>),
6682+
/// Replaces occurrences of `typing.Self` with a new `Self` type variable with the given upper bound.
6683+
ReplaceSelf { new_upper_bound: Type<'db> },
66656684
/// Marks the typevars that are bound by a generic class or function as inferable.
66666685
MarkTypeVarsInferable(BindingContext<'db>),
66676686
/// Create the top or bottom materialization of a type.
@@ -6683,6 +6702,9 @@ fn walk_type_mapping<'db, V: visitor::TypeVisitor<'db> + ?Sized>(
66836702
TypeMapping::BindSelf(self_type) => {
66846703
visitor.visit_type(db, *self_type);
66856704
}
6705+
TypeMapping::ReplaceSelf { new_upper_bound } => {
6706+
visitor.visit_type(db, *new_upper_bound);
6707+
}
66866708
TypeMapping::PromoteLiterals
66876709
| TypeMapping::BindLegacyTypevars(_)
66886710
| TypeMapping::MarkTypeVarsInferable(_)
@@ -6704,6 +6726,9 @@ impl<'db> TypeMapping<'_, 'db> {
67046726
TypeMapping::BindLegacyTypevars(*binding_context)
67056727
}
67066728
TypeMapping::BindSelf(self_type) => TypeMapping::BindSelf(*self_type),
6729+
TypeMapping::ReplaceSelf { new_upper_bound } => TypeMapping::ReplaceSelf {
6730+
new_upper_bound: *new_upper_bound,
6731+
},
67076732
TypeMapping::MarkTypeVarsInferable(binding_context) => {
67086733
TypeMapping::MarkTypeVarsInferable(*binding_context)
67096734
}
@@ -6728,6 +6753,9 @@ impl<'db> TypeMapping<'_, 'db> {
67286753
TypeMapping::BindSelf(self_type) => {
67296754
TypeMapping::BindSelf(self_type.normalized_impl(db, visitor))
67306755
}
6756+
TypeMapping::ReplaceSelf { new_upper_bound } => TypeMapping::ReplaceSelf {
6757+
new_upper_bound: new_upper_bound.normalized_impl(db, visitor),
6758+
},
67316759
TypeMapping::MarkTypeVarsInferable(binding_context) => {
67326760
TypeMapping::MarkTypeVarsInferable(*binding_context)
67336761
}
@@ -6736,6 +6764,37 @@ impl<'db> TypeMapping<'_, 'db> {
67366764
}
67376765
}
67386766
}
6767+
6768+
/// Update the generic context of a [`Signature`] according to the current type mapping
6769+
pub(crate) fn update_signature_generic_context(
6770+
&self,
6771+
db: &'db dyn Db,
6772+
context: GenericContext<'db>,
6773+
) -> GenericContext<'db> {
6774+
match self {
6775+
TypeMapping::Specialization(_)
6776+
| TypeMapping::PartialSpecialization(_)
6777+
| TypeMapping::PromoteLiterals
6778+
| TypeMapping::BindLegacyTypevars(_)
6779+
| TypeMapping::MarkTypeVarsInferable(_)
6780+
| TypeMapping::Materialize(_)
6781+
| TypeMapping::BindSelf(_) => context,
6782+
TypeMapping::ReplaceSelf { new_upper_bound } => GenericContext::from_typevar_instances(
6783+
db,
6784+
context.variables(db).iter().map(|typevar| {
6785+
if typevar.typevar(db).is_self(db) {
6786+
BoundTypeVarInstance::synthetic_self(
6787+
db,
6788+
*new_upper_bound,
6789+
typevar.binding_context(db),
6790+
)
6791+
} else {
6792+
*typevar
6793+
}
6794+
}),
6795+
),
6796+
}
6797+
}
67396798
}
67406799

67416800
/// A Salsa-tracked constraint set. This is only needed to have something appropriately small to
@@ -7663,6 +7722,27 @@ impl<'db> BoundTypeVarInstance<'db> {
76637722
)
76647723
}
76657724

7725+
/// Create a new synthetic `Self` type variable with the given upper bound.
7726+
pub(crate) fn synthetic_self(
7727+
db: &'db dyn Db,
7728+
upper_bound: Type<'db>,
7729+
binding_context: BindingContext<'db>,
7730+
) -> Self {
7731+
Self::new(
7732+
db,
7733+
TypeVarInstance::new(
7734+
db,
7735+
Name::new_static("Self"),
7736+
None,
7737+
Some(TypeVarBoundOrConstraints::UpperBound(upper_bound).into()),
7738+
Some(TypeVarVariance::Invariant),
7739+
None,
7740+
TypeVarKind::TypingSelf,
7741+
),
7742+
binding_context,
7743+
)
7744+
}
7745+
76667746
pub(crate) fn variance_with_polarity(
76677747
self,
76687748
db: &'db dyn Db,
@@ -10836,6 +10916,24 @@ impl<'db> TypeIsType<'db> {
1083610916
}
1083710917
}
1083810918

10919+
/// Walk the MRO of this class and return the last class just before the specified known base.
10920+
/// This can be used to determine upper bounds for `Self` type variables on methods that are
10921+
/// being added to the given class.
10922+
pub(super) fn determine_upper_bound<'db>(
10923+
db: &'db dyn Db,
10924+
class_literal: ClassLiteral<'db>,
10925+
specialization: Option<Specialization<'db>>,
10926+
is_known_base: impl Fn(ClassBase<'db>) -> bool,
10927+
) -> Type<'db> {
10928+
let upper_bound = class_literal
10929+
.iter_mro(db, specialization)
10930+
.take_while(|base| !is_known_base(*base))
10931+
.filter_map(ClassBase::into_class)
10932+
.last()
10933+
.unwrap_or_else(|| class_literal.unknown_specialization(db));
10934+
Type::instance(db, upper_bound)
10935+
}
10936+
1083910937
// Make sure that the `Type` enum does not grow unexpectedly.
1084010938
#[cfg(not(debug_assertions))]
1084110939
#[cfg(target_pointer_width = "64")]

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

Lines changed: 50 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,8 @@ use crate::types::{
3030
IsEquivalentVisitor, KnownInstanceType, ManualPEP695TypeAliasType, MaterializationKind,
3131
NormalizedVisitor, PropertyInstanceType, StringLiteralType, TypeAliasType, TypeContext,
3232
TypeMapping, TypeRelation, TypeVarBoundOrConstraints, TypeVarInstance, TypeVarKind,
33-
TypedDictParams, UnionBuilder, VarianceInferable, declaration_type, infer_definition_types,
33+
TypedDictParams, UnionBuilder, VarianceInferable, declaration_type, determine_upper_bound,
34+
infer_definition_types,
3435
};
3536
use crate::{
3637
Db, FxIndexMap, FxOrderSet, Program,
@@ -1975,7 +1976,20 @@ impl<'db> ClassLiteral<'db> {
19751976
return KnownClass::TypedDictFallback
19761977
.to_class_literal(db)
19771978
.find_name_in_mro_with_policy(db, name, policy)
1978-
.expect("Will return Some() when called on class literal");
1979+
.expect("Will return Some() when called on class literal")
1980+
.map_type(|ty| {
1981+
ty.apply_type_mapping(
1982+
db,
1983+
&TypeMapping::ReplaceSelf {
1984+
new_upper_bound: determine_upper_bound(
1985+
db,
1986+
self,
1987+
None,
1988+
ClassBase::is_typed_dict,
1989+
),
1990+
},
1991+
)
1992+
});
19791993
}
19801994
}
19811995
if lookup_result.is_ok() {
@@ -2256,6 +2270,22 @@ impl<'db> ClassLiteral<'db> {
22562270
.own_class_member(db, self.generic_context(db), None, name)
22572271
.place
22582272
.ignore_possibly_unbound()
2273+
.map(|ty| {
2274+
ty.apply_type_mapping(
2275+
db,
2276+
&TypeMapping::ReplaceSelf {
2277+
new_upper_bound: determine_upper_bound(
2278+
db,
2279+
self,
2280+
specialization,
2281+
|base| {
2282+
base.into_class()
2283+
.is_some_and(|c| c.is_known(db, KnownClass::Tuple))
2284+
},
2285+
),
2286+
},
2287+
)
2288+
})
22592289
}
22602290
(CodeGeneratorKind::DataclassLike, "__replace__")
22612291
if Program::get(db).python_version(db) >= PythonVersion::PY313 =>
@@ -2578,6 +2608,12 @@ impl<'db> ClassLiteral<'db> {
25782608
.to_class_literal(db)
25792609
.find_name_in_mro_with_policy(db, name, policy)
25802610
.expect("`find_name_in_mro_with_policy` will return `Some()` when called on class literal")
2611+
.map_type(|ty|
2612+
ty.apply_type_mapping(
2613+
db,
2614+
&TypeMapping::ReplaceSelf {new_upper_bound: determine_upper_bound(db, self, specialization, ClassBase::is_typed_dict) }
2615+
)
2616+
)
25812617
}
25822618
}
25832619

@@ -2815,7 +2851,18 @@ impl<'db> ClassLiteral<'db> {
28152851
ClassBase::TypedDict => {
28162852
return KnownClass::TypedDictFallback
28172853
.to_instance(db)
2818-
.instance_member(db, name);
2854+
.instance_member(db, name)
2855+
.map_type(|ty| {
2856+
ty.apply_type_mapping(
2857+
db,
2858+
&TypeMapping::ReplaceSelf {
2859+
new_upper_bound: Type::instance(
2860+
db,
2861+
self.unknown_specialization(db),
2862+
),
2863+
},
2864+
)
2865+
});
28192866
}
28202867
}
28212868
}

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,10 @@ impl<'db> ClassBase<'db> {
6969
.map_or(Self::unknown(), Self::Class)
7070
}
7171

72+
pub(super) const fn is_typed_dict(self) -> bool {
73+
matches!(self, ClassBase::TypedDict)
74+
}
75+
7276
/// Attempt to resolve `ty` into a `ClassBase`.
7377
///
7478
/// Return `None` if `ty` is not an acceptable type for a class base.

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -470,7 +470,9 @@ impl<'db> Signature<'db> {
470470
_ => type_mapping,
471471
};
472472
Self {
473-
generic_context: self.generic_context,
473+
generic_context: self
474+
.generic_context
475+
.map(|context| type_mapping.update_signature_generic_context(db, context)),
474476
inherited_generic_context: self.inherited_generic_context,
475477
definition: self.definition,
476478
parameters: self

0 commit comments

Comments
 (0)