[ty] Give Self type variables distinct class identities (astral-sh#29191)

Currently, all `Self` type variables share a single `TypeVarIdentity`.
This isn't really semantically right, but it's mostly fine in practice,
because they get bound to different methods and so have distinct
`BoundTypeVarIdentity`. It does cause bugs in some cases with
synthesized methods: unions of `_replace` methods from different named
tuple classes can collapse to a single callable because their synthetic
`Self` variables share an identity. This can incorrectly accept a
callable promised to return only one of the classes, depending on the
union order.

Also, we currently have to recover the owning class of a `Self` type
variable via its upper bound (which should be set to the owning class),
but this is a bit implicit and can be awkward when we transform upper
bounds (e.g. for `type[T]` or due to specialization or materialization).

Instead, give class-owned `Self` variables explicitly distinct
identities using `SelfTypeVarOrigin::Class`. Keep that origin stable
when specializing or materializing their bounds, and represent receiver
placeholders and synthesized `TypedDict` schemas with explicit unit
variants.

Author: carljm

crates/ty_ide/src/semantic_tokens.rs

@@ -441,7 +441,10 @@ impl<'db> SemanticTokenVisitor<'db> {
                     };
 
                     if let Some((type_var, is_cls)) = type_var
-                        && matches!(type_var.typevar(db).kind(db), TypeVarKind::TypingSelf)
+                        && matches!(
+                            type_var.typevar(db).kind(db),
+                            TypeVarKind::TypingSelf { .. }
+                        )
                     {
                         let kind = if is_cls {
                             SemanticTokenType::ClsParameter

crates/ty_python_semantic/resources/mdtest/named_tuple.md

@@ -1370,6 +1370,44 @@ class IntBox(Box[int]):
 reveal_type(IntBox(1)._replace(content=42))  # revealed: IntBox
 ```
 
+### `_replace` methods from different classes
+
+The `_replace` methods of different named tuples retain different return types even when their
+fields match. A union of these methods cannot promise to return only one of the classes, regardless
+of branch order:
+
+```py
+from typing import Callable, NamedTuple
+
+class A(NamedTuple):
+    x: int
+
+class B(NamedTuple):
+    x: int
+
+def choose(flag: bool) -> Callable[..., A]:
+    # error: [invalid-return-type]
+    return A._replace if flag else B._replace
+
+def reversed_order(flag: bool) -> Callable[..., A]:
+    # error: [invalid-return-type]
+    return B._replace if flag else A._replace
+
+def either_class(flag: bool) -> Callable[..., A | B]:
+    return A._replace if flag else B._replace  # no diagnostic
+```
+
+The same distinction applies to named tuples created with the functional syntax:
+
+```py
+DynamicA = NamedTuple("DynamicA", [("x", int)])
+DynamicB = NamedTuple("DynamicB", [("x", int)])
+
+def choose_dynamic(flag: bool) -> Callable[..., DynamicA]:
+    # error: [invalid-return-type]
+    return DynamicA._replace if flag else DynamicB._replace
+```
+
 ## `collections.namedtuple`
 
 ```py

crates/ty_python_semantic/src/types.rs

@@ -106,8 +106,8 @@ pub use crate::types::type_alias::TypeAliasType;
 pub use crate::types::type_form::TypeFormType;
 pub(crate) use crate::types::typed_dict::TypedDictType;
 pub(crate) use crate::types::typevar::{
-    BindingContext, BoundTypeVarIdentity, ParamSpecAttrKind, TypeVarBoundOrConstraints,
-    TypeVarNonce,
+    BindingContext, BoundTypeVarIdentity, ParamSpecAttrKind, SelfTypeVarOrigin,
+    TypeVarBoundOrConstraints, TypeVarNonce,
 };
 pub use crate::types::typevar::{BoundTypeVarInstance, TypeVarKind};
 use crate::types::typevar::{TypeVarInstance, TypeVarSet};
@@ -10000,7 +10000,9 @@ impl<'db> Type<'db> {
     ) {
         let matching_typevar = |bound_typevar: &BoundTypeVarInstance<'db>| {
             match bound_typevar.typevar(db).kind(db) {
-                TypeVarKind::LegacyTypeVar | TypeVarKind::Pep613Alias | TypeVarKind::TypingSelf
+                TypeVarKind::LegacyTypeVar
+                | TypeVarKind::Pep613Alias
+                | TypeVarKind::TypingSelf { .. }
                     if binding_context.is_none_or(|binding_context| {
                         bound_typevar.binding_context(db)
                             == BindingContext::Definition(binding_context)
@@ -11146,17 +11148,17 @@ pub enum PromotionKind {
     SingletonsOnly,
 }
 
-/// Returns the [`ClassLiteral`] that "owns" a `Self` typevar (i.e., the class from its upper bound).
-fn self_typevar_owner_class_literal<'db>(
+/// Returns the class recorded in a `Self` typevar's origin.
+fn self_typevar_origin_class_literal<'db>(
     db: &'db dyn Db,
-    env: &ProgramEnvironment<'db>,
     bound_typevar: BoundTypeVarInstance<'db>,
 ) -> Option<ClassLiteral<'db>> {
-    bound_typevar
-        .typevar(db)
-        .upper_bound(db, env)
-        .and_then(|ty| ty.nominal_class(db, env))
-        .map(|class| class.class_literal(db))
+    match bound_typevar.kind(db) {
+        TypeVarKind::TypingSelf {
+            origin: SelfTypeVarOrigin::Class(class),
+        } => Some(class),
+        _ => None,
+    }
 }
 
 #[salsa::tracked(returns(ref), heap_size=ruff_memory_usage::heap_size)]
@@ -11173,8 +11175,8 @@ fn class_mro_literals<'db>(
 
 /// Information needed to bind `Self` typevars to a concrete type.
 ///
-/// Uses MRO-based matching: a `Self` typevar is bound only if its owner class
-/// is in the MRO of the self type's class.
+/// Uses the binding context and MRO-based matching against the class recorded in a `Self`
+/// typevar's origin.
 #[derive(Clone, Debug, Eq, PartialEq, get_size2::GetSize)]
 pub struct SelfBinding<'db> {
     ty: Type<'db>,
@@ -11201,7 +11203,7 @@ impl<'db> SelfBinding<'db> {
     ) -> Self {
         let class_literal = match self_type {
             Type::TypeVar(typevar) if typevar.typevar(db).is_self(db) => {
-                self_typevar_owner_class_literal(db, env, typevar)
+                self_typevar_origin_class_literal(db, typevar)
             }
             _ => self_type
                 .nominal_class(db, env)
@@ -11216,12 +11218,7 @@ impl<'db> SelfBinding<'db> {
     }
 
     /// Returns whether `bound_typevar` should be replaced by this binding's concrete self type.
-    fn should_bind(
-        &self,
-        db: &'db dyn Db,
-        env: &ProgramEnvironment<'db>,
-        bound_typevar: BoundTypeVarInstance<'db>,
-    ) -> bool {
+    fn should_bind(&self, db: &'db dyn Db, bound_typevar: BoundTypeVarInstance<'db>) -> bool {
         if !bound_typevar.typevar(db).is_self(db) {
             return false;
         }
@@ -11232,12 +11229,12 @@ impl<'db> SelfBinding<'db> {
             return true;
         }
 
-        // Check that the Self typevar's owner class is in the MRO of the self type's class.
-        // If we can't determine either class, conservatively don't bind.
+        // Check that Self's originating class is in the receiver's MRO. Receiver placeholders
+        // and synthesized TypedDict schemas have no class to restrict matching.
         self.class_literal.is_some_and(|class_literal| {
             let class_mro = class_mro_literals(db, class_literal);
-            self_typevar_owner_class_literal(db, env, bound_typevar)
-                .is_none_or(|owner_class| class_mro.contains(&owner_class))
+            self_typevar_origin_class_literal(db, bound_typevar)
+                .is_none_or(|origin_class| class_mro.contains(&origin_class))
         })
     }
 }
@@ -11274,8 +11271,11 @@ pub enum TypeMapping<'a, 'db> {
     },
     /// Binds any `typing.Self` typevar with a particular `self` class.
     BindSelf(SelfBinding<'db>),
-    /// Replaces occurrences of `typing.Self` with a new `Self` type variable with the given upper bound.
-    ReplaceSelf { new_upper_bound: Type<'db> },
+    /// Replaces `typing.Self` when adapting a synthesized method to another class or schema.
+    ReplaceSelf {
+        new_origin: SelfTypeVarOrigin<'db>,
+        new_upper_bound: Type<'db>,
+    },
     /// Create the top or bottom materialization of a type.
     Materialize(MaterializationKind),
     /// Replace default types in parameters of callables with `Unknown`. This is used to avoid infinite
@@ -11358,23 +11358,25 @@ impl<'db> TypeMapping<'_, 'db> {
                     Some(context)
                 }
             }
-            TypeMapping::ReplaceSelf { new_upper_bound } => {
-                Some(GenericContext::from_typevar_instances(
-                    db,
-                    env,
-                    context.variables(db).map(|typevar| {
-                        if typevar.typevar(db).is_self(db) {
-                            BoundTypeVarInstance::synthetic_self(
-                                db,
-                                *new_upper_bound,
-                                typevar.binding_context(db),
-                            )
-                        } else {
-                            typevar
-                        }
-                    }),
-                ))
-            }
+            TypeMapping::ReplaceSelf {
+                new_origin,
+                new_upper_bound,
+            } => Some(GenericContext::from_typevar_instances(
+                db,
+                env,
+                context.variables(db).map(|typevar| {
+                    if typevar.typevar(db).is_self(db) {
+                        BoundTypeVarInstance::synthetic_self(
+                            db,
+                            *new_origin,
+                            *new_upper_bound,
+                            typevar.binding_context(db),
+                        )
+                    } else {
+                        typevar
+                    }
+                }),
+            )),
         }
     }
 
@@ -12652,24 +12654,21 @@ impl<'db> TypeGuardLike<'db> for TypeGuardType<'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.
+/// This determines the owner and upper bound of `Self` on synthesized fallback methods.
 ///
 /// Preserve the class's specialization so that a method on `Child[int]` has a bound such as
 /// `Base[int]`, rather than retaining the type variable in `Base[T@Child]`.
-pub(super) fn determine_upper_bound<'db>(
+pub(super) fn determine_self_class<'db>(
     db: &'db dyn Db,
-    env: &ProgramEnvironment<'db>,
     class: ClassType<'db>,
     is_known_base: impl Fn(ClassBase<'db>) -> bool,
-) -> Type<'db> {
-    let upper_bound = class
+) -> ClassType<'db> {
+    class
         .iter_mro(db)
         .take_while(|base| !is_known_base(*base))
         .filter_map(ClassBase::into_class)
         .last()
-        .unwrap_or(class);
-    Type::instance(db, env, upper_bound)
+        .unwrap_or(class)
 }
 
 // Make sure that the `Type` enum does not grow unexpectedly.

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

@@ -9,7 +9,8 @@ use crate::{
     types::{
         BindingContext, BoundTypeVarInstance, ClassBase, ClassLiteral, ClassType, GenericContext,
         KnownClass, KnownInstanceType, MemberLookupPolicy, Parameter, Parameters,
-        PropertyInstanceType, Signature, SubclassOfType, Type, TypeContext, TypeMapping,
+        PropertyInstanceType, SelfTypeVarOrigin, Signature, SubclassOfType, Type, TypeContext,
+        TypeMapping,
         class::{DynamicClassHeaderAnchor, DynamicClassScopeOffset, dynamic_class_header_range},
         definition_expression_type,
         member::Member,
@@ -30,15 +31,18 @@ pub(super) fn synthesize_namedtuple_class_member<'db>(
     db: &'db dyn Db,
     env: &ProgramEnvironment<'db>,
     name: &str,
-    instance_ty: Type<'db>,
+    class: ClassType<'db>,
     fields: impl Iterator<Item = NamedTupleField<'db>>,
     inherited_generic_context: Option<GenericContext<'db>>,
 ) -> Option<Type<'db>> {
+    let instance_ty = Type::instance(db, env, class);
+    let origin = SelfTypeVarOrigin::Class(class.class_literal(db));
     match name {
         "__new__" => {
             // __new__(cls, field1, field2, ...) -> Self
             let self_typevar = BoundTypeVarInstance::synthetic_self(
                 db,
+                origin,
                 instance_ty,
                 BindingContext::Synthetic(env.program(db)),
             );
@@ -96,6 +100,7 @@ pub(super) fn synthesize_namedtuple_class_member<'db>(
             // _replace(self, *, field1=..., field2=...) -> Self
             let self_ty = Type::TypeVar(BoundTypeVarInstance::synthetic_self(
                 db,
+                origin,
                 instance_ty,
                 BindingContext::Synthetic(env.program(db)),
             ));
@@ -398,6 +403,7 @@ impl<'db> DynamicNamedTupleLiteral<'db> {
                                 db,
                                 env,
                                 &TypeMapping::ReplaceSelf {
+                                    new_origin: SelfTypeVarOrigin::Class(self.into()),
                                     new_upper_bound: instance_ty,
                                 },
                                 TypeContext::default(),
@@ -412,7 +418,7 @@ impl<'db> DynamicNamedTupleLiteral<'db> {
             db,
             env,
             name,
-            instance_ty,
+            ClassType::NonGeneric(self.into()),
             self.fields(db).iter().cloned(),
             None,
         );
@@ -430,6 +436,7 @@ impl<'db> DynamicNamedTupleLiteral<'db> {
                     db,
                     env,
                     &TypeMapping::ReplaceSelf {
+                        new_origin: SelfTypeVarOrigin::Class(self.into()),
                         new_upper_bound: instance_ty,
                     },
                     TypeContext::default(),

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

@@ -25,9 +25,9 @@ use crate::{
         CallableType, ClassBase, ClassLiteral, ClassType, DATACLASS_FLAGS, DataclassFlags,
         DataclassParams, GenericAlias, GenericContext, KnownClass, KnownInstanceType,
         MaterializationKind, MemberLookupPolicy, MetaclassCandidate, MetaclassTransformInfo,
-        Parameter, Parameters, PropertyInstanceType, Signature, SpecialFormType, StaticMroError,
-        SubclassOfType, Type, TypeContext, TypeMapping, TypeVarVariance, TypingModule,
-        UnionBuilder, UnionType,
+        Parameter, Parameters, PropertyInstanceType, SelfTypeVarOrigin, Signature, SpecialFormType,
+        StaticMroError, SubclassOfType, Type, TypeContext, TypeMapping, TypeVarVariance,
+        TypingModule, UnionBuilder, UnionType,
         attribute_write::DescriptorSetterDomain,
         bound_super::BoundSuperType,
         call::{CallError, CallErrorKind},
@@ -41,7 +41,7 @@ use crate::{
         },
         context::InferContext,
         dedicated::pydantic,
-        definition_expression_type, determine_upper_bound,
+        definition_expression_type, determine_self_class,
         diagnostic::INVALID_DATACLASS_OVERRIDE,
         enums::{enum_metadata, is_enum_class_by_inheritance, try_unwrap_nonmember_value},
         function::{DataclassTransformerParams, KnownFunction},
@@ -2031,6 +2031,7 @@ impl<'db> StaticClassLiteral<'db> {
                             db,
                             env,
                             &TypeMapping::ReplaceSelf {
+                                new_origin: SelfTypeVarOrigin::Class(self.into()),
                                 new_upper_bound: instance_ty,
                             },
                             TypeContext::default(),
@@ -2058,7 +2059,7 @@ impl<'db> StaticClassLiteral<'db> {
                     db,
                     env,
                     name,
-                    instance_ty,
+                    self.apply_optional_specialization(db, specialization),
                     fields_iter,
                     specialization.map(|s| s.generic_context(db)),
                 )
@@ -2139,19 +2140,20 @@ impl<'db> StaticClassLiteral<'db> {
                     .own_class_member(db, env, self.inherited_generic_context(db), None, name)
                     .ignore_possibly_undefined()
                     .map(|ty| {
+                        let self_class = determine_self_class(
+                            db,
+                            self.apply_optional_specialization(db, specialization),
+                            |base| {
+                                base.into_class()
+                                    .is_some_and(|c| c.is_known(db, KnownClass::Tuple))
+                            },
+                        );
                         ty.apply_type_mapping(
                             db,
                             env,
                             &TypeMapping::ReplaceSelf {
-                                new_upper_bound: determine_upper_bound(
-                                    db,
-                                    env,
-                                    self.apply_optional_specialization(db, specialization),
-                                    |base| {
-                                        base.into_class()
-                                            .is_some_and(|c| c.is_known(db, KnownClass::Tuple))
-                                    },
-                                ),
+                                new_origin: SelfTypeVarOrigin::Class(self_class.class_literal(db)),
+                                new_upper_bound: Type::instance(db, env, self_class),
                             },
                             TypeContext::default(),
                         )
@@ -2940,6 +2942,7 @@ impl<'db> StaticClassLiteral<'db> {
                         db,
                         env,
                         &TypeMapping::ReplaceSelf {
+                            new_origin: SelfTypeVarOrigin::Class(self.into()),
                             new_upper_bound: Type::instance(
                                 db,
                                 env,