[ty] Distribute type[] over unions (#22115)

## Summary

Closes https://github.com/astral-sh/ty/issues/2121.
This commit is contained in:
Charlie Marsh
2025-12-21 18:45:29 -05:00
committed by GitHub
parent b6e84eca16
commit fee4e2d72a
4 changed files with 94 additions and 6 deletions

View File

@@ -7743,11 +7743,7 @@ impl<'db> Type<'db> {
}
Type::ClassLiteral(class) => class.metaclass(db),
Type::GenericAlias(alias) => ClassType::from(alias).metaclass(db),
Type::SubclassOf(subclass_of_ty) => match subclass_of_ty.subclass_of().into_class(db) {
None => self,
Some(class) => SubclassOfType::try_from_type(db, class.metaclass(db))
.unwrap_or(SubclassOfType::subclass_of_unknown()),
},
Type::SubclassOf(subclass_of_ty) => subclass_of_ty.to_meta_type(db),
Type::StringLiteral(_) | Type::LiteralString => KnownClass::Str.to_class_literal(db),
Type::Dynamic(dynamic) => SubclassOfType::from(db, SubclassOfInner::Dynamic(dynamic)),
// TODO intersections

View File

@@ -8,7 +8,7 @@ use crate::types::{
ApplyTypeMappingVisitor, BoundTypeVarInstance, ClassType, DynamicType,
FindLegacyTypeVarsVisitor, HasRelationToVisitor, IsDisjointVisitor, KnownClass,
MaterializationKind, MemberLookupPolicy, NormalizedVisitor, SpecialFormType, Type, TypeContext,
TypeMapping, TypeRelation, TypeVarBoundOrConstraints, TypedDictType, todo_type,
TypeMapping, TypeRelation, TypeVarBoundOrConstraints, TypedDictType, UnionType, todo_type,
};
use crate::{Db, FxOrderSet};
@@ -79,6 +79,18 @@ impl<'db> SubclassOfType<'db> {
/// Given an instance of the class or type variable `T`, returns a [`Type`] instance representing `type[T]`.
pub(crate) fn try_from_instance(db: &'db dyn Db, ty: Type<'db>) -> Option<Type<'db>> {
// Handle unions by distributing `type[]` over each element:
// `type[A | B]` -> `type[A] | type[B]`
if let Type::Union(union) = ty {
return UnionType::try_from_elements(
db,
union
.elements(db)
.iter()
.map(|element| Self::try_from_instance(db, *element)),
);
}
SubclassOfInner::try_from_instance(db, ty).map(|subclass_of| Self::from(db, subclass_of))
}
@@ -290,6 +302,35 @@ impl<'db> SubclassOfType<'db> {
}
}
/// Compute the metatype of this `type[T]`.
///
/// For `type[C]` where `C` is a concrete class, this returns `type[metaclass(C)]`.
/// For `type[T]` where `T` is a `TypeVar`, this computes the metatype based on the
/// `TypeVar`'s bounds or constraints.
pub(crate) fn to_meta_type(self, db: &'db dyn Db) -> Type<'db> {
match self.subclass_of.with_transposed_type_var(db) {
SubclassOfInner::Dynamic(dynamic) => {
SubclassOfType::from(db, SubclassOfInner::Dynamic(dynamic))
}
SubclassOfInner::Class(class) => SubclassOfType::try_from_type(db, class.metaclass(db))
.unwrap_or(SubclassOfType::subclass_of_unknown()),
// For `type[T]` where `T` is a TypeVar, `with_transposed_type_var` transforms
// the bounds from instance types to `type[]` types. For example, `type[T]` where
// `T: A | B` becomes a TypeVar with bound `type[A] | type[B]`. The metatype is
// then the metatype of that bound.
SubclassOfInner::TypeVar(bound_typevar) => {
match bound_typevar.typevar(db).bound_or_constraints(db) {
// `with_transposed_type_var` always adds a bound for unbounded TypeVars
None => unreachable!(),
Some(TypeVarBoundOrConstraints::UpperBound(bound)) => bound.to_meta_type(db),
Some(TypeVarBoundOrConstraints::Constraints(constraints)) => {
constraints.as_type(db).to_meta_type(db)
}
}
}
}
}
pub(crate) fn is_typed_dict(self, db: &'db dyn Db) -> bool {
self.subclass_of
.into_class(db)