[ty] Patch Self for fallback-methods on NamedTuples and TypedDicts (#20328)

## Summary

We use classes like
[`_typeshed._type_checker_internals.NamedTupleFallback`](d9c76e1d9f/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 feb22aae00..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>

crates/ty_python_semantic/resources/mdtest/named_tuple.md

@@ -272,16 +272,34 @@ reveal_type(Person._make)  # revealed: bound method <class 'Person'>._make(itera
 reveal_type(Person._asdict)  # revealed: def _asdict(self) -> dict[str, Any]
 reveal_type(Person._replace)  # revealed: def _replace(self, **kwargs: Any) -> Self@_replace
 
-# TODO: should be `Person` once we support `Self`
+# TODO: should be `Person` once we support implicit type of `self`
 reveal_type(Person._make(("Alice", 42)))  # revealed: Unknown
 
 person = Person("Alice", 42)
 
 reveal_type(person._asdict())  # revealed: dict[str, Any]
-# TODO: should be `Person` once we support `Self`
+# TODO: should be `Person` once we support implicit type of `self`
 reveal_type(person._replace(name="Bob"))  # revealed: Unknown
 ```
 
+When accessing them on child classes of generic `NamedTuple`s, the return type is specialized
+accordingly:
+
+```py
+from typing import NamedTuple, Generic, TypeVar
+
+T = TypeVar("T")
+
+class Box(NamedTuple, Generic[T]):
+    content: T
+
+class IntBox(Box[int]):
+    pass
+
+# TODO: should be `IntBox` once we support the implicit type of `self`
+reveal_type(IntBox(1)._replace(content=42))  # revealed: Unknown
+```
+
 ## `collections.namedtuple`
 
 ```py

crates/ty_python_semantic/resources/mdtest/typed_dict.md

@@ -627,6 +627,18 @@ alice: Employee = {"name": "Alice", "employee_id": 1}
 
 # error: [missing-typed-dict-key] "Missing required key 'employee_id' in TypedDict `Employee` constructor"
 eve: Employee = {"name": "Eve"}
+
+def combine(p: Person, e: Employee):
+    # TODO: Should be `Person` once we support the implicit type of self
+    reveal_type(p.copy())  # revealed: Unknown
+    # TODO: Should be `Employee` once we support the implicit type of self
+    reveal_type(e.copy())  # revealed: Unknown
+
+    reveal_type(p | p)  # revealed: Person
+    reveal_type(e | e)  # revealed: Employee
+
+    # TODO: Should be `Person` once we support the implicit type of self and subtyping for TypedDicts
+    reveal_type(p | e)  # revealed: Employee
 ```
 
 When inheriting from a `TypedDict` with a different `total` setting, inherited fields maintain their

crates/ty_python_semantic/src/types.rs

@@ -5971,6 +5971,19 @@ impl<'db> Type<'db> {
                         self
                     }
                 }
+                TypeMapping::ReplaceSelf { new_upper_bound } => {
+                    if bound_typevar.typevar(db).is_self(db) {
+                        Type::TypeVar(
+                            BoundTypeVarInstance::synthetic_self(
+                                db,
+                                *new_upper_bound,
+                                bound_typevar.binding_context(db)
+                            )
+                        )
+                    } else {
+                        self
+                    }
+                }
                 TypeMapping::PromoteLiterals | TypeMapping::BindLegacyTypevars(_) |
                 TypeMapping::MarkTypeVarsInferable(_) => self,
                 TypeMapping::Materialize(materialization_kind)  => {
@@ -5994,7 +6007,8 @@ impl<'db> Type<'db> {
                 }
                 TypeMapping::PromoteLiterals |
                 TypeMapping::BindLegacyTypevars(_) |
-                TypeMapping::BindSelf(_)
+                TypeMapping::BindSelf(_) |
+                TypeMapping::ReplaceSelf { .. }
                     => self,
                 TypeMapping::Materialize(materialization_kind)  => Type::NonInferableTypeVar(bound_typevar.materialize_impl(db, *materialization_kind, visitor))
 
@@ -6008,6 +6022,7 @@ impl<'db> Type<'db> {
                 TypeMapping::PartialSpecialization(_) |
                 TypeMapping::PromoteLiterals |
                 TypeMapping::BindSelf(_) |
+                TypeMapping::ReplaceSelf { .. } |
                 TypeMapping::MarkTypeVarsInferable(_) |
                 TypeMapping::Materialize(_) => self,
             }
@@ -6116,6 +6131,7 @@ impl<'db> Type<'db> {
                 TypeMapping::PartialSpecialization(_) |
                 TypeMapping::BindLegacyTypevars(_) |
                 TypeMapping::BindSelf(_) |
+                TypeMapping::ReplaceSelf { .. } |
                 TypeMapping::MarkTypeVarsInferable(_) |
                 TypeMapping::Materialize(_) => self,
                 TypeMapping::PromoteLiterals => self.literal_fallback_instance(db)
@@ -6127,6 +6143,7 @@ impl<'db> Type<'db> {
                 TypeMapping::PartialSpecialization(_) |
                 TypeMapping::BindLegacyTypevars(_) |
                 TypeMapping::BindSelf(_) |
+                TypeMapping::ReplaceSelf { .. } |
                 TypeMapping::MarkTypeVarsInferable(_) |
                 TypeMapping::PromoteLiterals => self,
                 TypeMapping::Materialize(materialization_kind) => match materialization_kind {
@@ -6662,6 +6679,8 @@ pub enum TypeMapping<'a, 'db> {
     BindLegacyTypevars(BindingContext<'db>),
     /// Binds any `typing.Self` typevar with a particular `self` class.
     BindSelf(Type<'db>),
+    /// Replaces occurrences of `typing.Self` with a new `Self` type variable with the given upper bound.
+    ReplaceSelf { new_upper_bound: Type<'db> },
     /// Marks the typevars that are bound by a generic class or function as inferable.
     MarkTypeVarsInferable(BindingContext<'db>),
     /// Create the top or bottom materialization of a type.
@@ -6683,6 +6702,9 @@ fn walk_type_mapping<'db, V: visitor::TypeVisitor<'db> + ?Sized>(
         TypeMapping::BindSelf(self_type) => {
             visitor.visit_type(db, *self_type);
         }
+        TypeMapping::ReplaceSelf { new_upper_bound } => {
+            visitor.visit_type(db, *new_upper_bound);
+        }
         TypeMapping::PromoteLiterals
         | TypeMapping::BindLegacyTypevars(_)
         | TypeMapping::MarkTypeVarsInferable(_)
@@ -6704,6 +6726,9 @@ impl<'db> TypeMapping<'_, 'db> {
                 TypeMapping::BindLegacyTypevars(*binding_context)
             }
             TypeMapping::BindSelf(self_type) => TypeMapping::BindSelf(*self_type),
+            TypeMapping::ReplaceSelf { new_upper_bound } => TypeMapping::ReplaceSelf {
+                new_upper_bound: *new_upper_bound,
+            },
             TypeMapping::MarkTypeVarsInferable(binding_context) => {
                 TypeMapping::MarkTypeVarsInferable(*binding_context)
             }
@@ -6728,6 +6753,9 @@ impl<'db> TypeMapping<'_, 'db> {
             TypeMapping::BindSelf(self_type) => {
                 TypeMapping::BindSelf(self_type.normalized_impl(db, visitor))
             }
+            TypeMapping::ReplaceSelf { new_upper_bound } => TypeMapping::ReplaceSelf {
+                new_upper_bound: new_upper_bound.normalized_impl(db, visitor),
+            },
             TypeMapping::MarkTypeVarsInferable(binding_context) => {
                 TypeMapping::MarkTypeVarsInferable(*binding_context)
             }
@@ -6736,6 +6764,37 @@ impl<'db> TypeMapping<'_, 'db> {
             }
         }
     }
+
+    /// Update the generic context of a [`Signature`] according to the current type mapping
+    pub(crate) fn update_signature_generic_context(
+        &self,
+        db: &'db dyn Db,
+        context: GenericContext<'db>,
+    ) -> GenericContext<'db> {
+        match self {
+            TypeMapping::Specialization(_)
+            | TypeMapping::PartialSpecialization(_)
+            | TypeMapping::PromoteLiterals
+            | TypeMapping::BindLegacyTypevars(_)
+            | TypeMapping::MarkTypeVarsInferable(_)
+            | TypeMapping::Materialize(_)
+            | TypeMapping::BindSelf(_) => context,
+            TypeMapping::ReplaceSelf { new_upper_bound } => GenericContext::from_typevar_instances(
+                db,
+                context.variables(db).iter().map(|typevar| {
+                    if typevar.typevar(db).is_self(db) {
+                        BoundTypeVarInstance::synthetic_self(
+                            db,
+                            *new_upper_bound,
+                            typevar.binding_context(db),
+                        )
+                    } else {
+                        *typevar
+                    }
+                }),
+            ),
+        }
+    }
 }
 
 /// A Salsa-tracked constraint set. This is only needed to have something appropriately small to
@@ -7663,6 +7722,27 @@ impl<'db> BoundTypeVarInstance<'db> {
         )
     }
 
+    /// Create a new synthetic `Self` type variable with the given upper bound.
+    pub(crate) fn synthetic_self(
+        db: &'db dyn Db,
+        upper_bound: Type<'db>,
+        binding_context: BindingContext<'db>,
+    ) -> Self {
+        Self::new(
+            db,
+            TypeVarInstance::new(
+                db,
+                Name::new_static("Self"),
+                None,
+                Some(TypeVarBoundOrConstraints::UpperBound(upper_bound).into()),
+                Some(TypeVarVariance::Invariant),
+                None,
+                TypeVarKind::TypingSelf,
+            ),
+            binding_context,
+        )
+    }
+
     pub(crate) fn variance_with_polarity(
         self,
         db: &'db dyn Db,
@@ -10836,6 +10916,24 @@ impl<'db> TypeIsType<'db> {
     }
 }
 
+/// Walk the MRO of this class and return the last class just before the specified known base.
+/// This can be used to determine upper bounds for `Self` type variables on methods that are
+/// being added to the given class.
+pub(super) fn determine_upper_bound<'db>(
+    db: &'db dyn Db,
+    class_literal: ClassLiteral<'db>,
+    specialization: Option<Specialization<'db>>,
+    is_known_base: impl Fn(ClassBase<'db>) -> bool,
+) -> Type<'db> {
+    let upper_bound = class_literal
+        .iter_mro(db, specialization)
+        .take_while(|base| !is_known_base(*base))
+        .filter_map(ClassBase::into_class)
+        .last()
+        .unwrap_or_else(|| class_literal.unknown_specialization(db));
+    Type::instance(db, upper_bound)
+}
+
 // Make sure that the `Type` enum does not grow unexpectedly.
 #[cfg(not(debug_assertions))]
 #[cfg(target_pointer_width = "64")]

crates/ty_python_semantic/src/types/class.rs

@@ -30,7 +30,8 @@ use crate::types::{
     IsEquivalentVisitor, KnownInstanceType, ManualPEP695TypeAliasType, MaterializationKind,
     NormalizedVisitor, PropertyInstanceType, StringLiteralType, TypeAliasType, TypeContext,
     TypeMapping, TypeRelation, TypeVarBoundOrConstraints, TypeVarInstance, TypeVarKind,
-    TypedDictParams, UnionBuilder, VarianceInferable, declaration_type, infer_definition_types,
+    TypedDictParams, UnionBuilder, VarianceInferable, declaration_type, determine_upper_bound,
+    infer_definition_types,
 };
 use crate::{
     Db, FxIndexMap, FxOrderSet, Program,
@@ -1975,7 +1976,20 @@ impl<'db> ClassLiteral<'db> {
                     return KnownClass::TypedDictFallback
                         .to_class_literal(db)
                         .find_name_in_mro_with_policy(db, name, policy)
-                        .expect("Will return Some() when called on class literal");
+                        .expect("Will return Some() when called on class literal")
+                        .map_type(|ty| {
+                            ty.apply_type_mapping(
+                                db,
+                                &TypeMapping::ReplaceSelf {
+                                    new_upper_bound: determine_upper_bound(
+                                        db,
+                                        self,
+                                        None,
+                                        ClassBase::is_typed_dict,
+                                    ),
+                                },
+                            )
+                        });
                 }
             }
             if lookup_result.is_ok() {
@@ -2256,6 +2270,22 @@ impl<'db> ClassLiteral<'db> {
                     .own_class_member(db, self.generic_context(db), None, name)
                     .place
                     .ignore_possibly_unbound()
+                    .map(|ty| {
+                        ty.apply_type_mapping(
+                            db,
+                            &TypeMapping::ReplaceSelf {
+                                new_upper_bound: determine_upper_bound(
+                                    db,
+                                    self,
+                                    specialization,
+                                    |base| {
+                                        base.into_class()
+                                            .is_some_and(|c| c.is_known(db, KnownClass::Tuple))
+                                    },
+                                ),
+                            },
+                        )
+                    })
             }
             (CodeGeneratorKind::DataclassLike, "__replace__")
                 if Program::get(db).python_version(db) >= PythonVersion::PY313 =>
@@ -2578,6 +2608,12 @@ impl<'db> ClassLiteral<'db> {
                 .to_class_literal(db)
                 .find_name_in_mro_with_policy(db, name, policy)
                 .expect("`find_name_in_mro_with_policy` will return `Some()` when called on class literal")
+                .map_type(|ty|
+                    ty.apply_type_mapping(
+                        db,
+                        &TypeMapping::ReplaceSelf {new_upper_bound: determine_upper_bound(db, self, specialization, ClassBase::is_typed_dict) }
+                    )
+                )
         }
     }
 
@@ -2815,7 +2851,18 @@ impl<'db> ClassLiteral<'db> {
                 ClassBase::TypedDict => {
                     return KnownClass::TypedDictFallback
                         .to_instance(db)
-                        .instance_member(db, name);
+                        .instance_member(db, name)
+                        .map_type(|ty| {
+                            ty.apply_type_mapping(
+                                db,
+                                &TypeMapping::ReplaceSelf {
+                                    new_upper_bound: Type::instance(
+                                        db,
+                                        self.unknown_specialization(db),
+                                    ),
+                                },
+                            )
+                        });
                 }
             }
         }

crates/ty_python_semantic/src/types/class_base.rs

@@ -69,6 +69,10 @@ impl<'db> ClassBase<'db> {
             .map_or(Self::unknown(), Self::Class)
     }
 
+    pub(super) const fn is_typed_dict(self) -> bool {
+        matches!(self, ClassBase::TypedDict)
+    }
+
     /// Attempt to resolve `ty` into a `ClassBase`.
     ///
     /// Return `None` if `ty` is not an acceptable type for a class base.

crates/ty_python_semantic/src/types/signatures.rs

@@ -470,7 +470,9 @@ impl<'db> Signature<'db> {
             _ => type_mapping,
         };
         Self {
-            generic_context: self.generic_context,
+            generic_context: self
+                .generic_context
+                .map(|context| type_mapping.update_signature_generic_context(db, context)),
             inherited_generic_context: self.inherited_generic_context,
             definition: self.definition,
             parameters: self