[ty] implement TypedDict structural assignment (#21467)

Closes https://github.com/astral-sh/ty/issues/1387.
This commit is contained in:
Jack O'Connor
2025-11-20 13:15:28 -08:00
committed by GitHub
parent 1b28fc1f14
commit eb7c098d6b
11 changed files with 678 additions and 59 deletions

View File

@@ -2093,14 +2093,6 @@ impl<'db> Type<'db> {
ConstraintSet::from(false)
}
(Type::TypedDict(_), _) => {
// TODO: Implement assignability and subtyping for TypedDict
ConstraintSet::from(relation.is_assignability())
}
// A non-`TypedDict` cannot subtype a `TypedDict`
(_, Type::TypedDict(_)) => ConstraintSet::from(false),
// Note that the definition of `Type::AlwaysFalsy` depends on the return value of `__bool__`.
// If `__bool__` always returns True or False, it can be treated as a subtype of `AlwaysTruthy` or `AlwaysFalsy`, respectively.
(left, Type::AlwaysFalsy) => ConstraintSet::from(left.bool(db).is_always_false()),
@@ -2207,6 +2199,38 @@ impl<'db> Type<'db> {
// A protocol instance can never be a subtype of a nominal type, with the *sole* exception of `object`.
(Type::ProtocolInstance(_), _) => ConstraintSet::from(false),
(Type::TypedDict(self_typeddict), Type::TypedDict(other_typeddict)) => relation_visitor
.visit((self, target, relation), || {
self_typeddict.has_relation_to_impl(
db,
other_typeddict,
inferable,
relation,
relation_visitor,
disjointness_visitor,
)
}),
// TODO: When we support `closed` and/or `extra_items`, we could allow assignments to other
// compatible `Mapping`s. `extra_items` could also allow for some assignments to `dict`, as
// long as `total=False`. (But then again, does anyone want a non-total `TypedDict` where all
// key types are a supertype of the extra items type?)
(Type::TypedDict(_), _) => relation_visitor.visit((self, target, relation), || {
KnownClass::Mapping
.to_specialized_instance(db, [KnownClass::Str.to_instance(db), Type::object()])
.has_relation_to_impl(
db,
target,
inferable,
relation,
relation_visitor,
disjointness_visitor,
)
}),
// A non-`TypedDict` cannot subtype a `TypedDict`
(_, Type::TypedDict(_)) => ConstraintSet::from(false),
// All `StringLiteral` types are a subtype of `LiteralString`.
(Type::StringLiteral(_), Type::LiteralString) => ConstraintSet::from(true),

View File

@@ -502,6 +502,17 @@ impl<'db> UnionBuilder<'db> {
return;
}
// Comparing `TypedDict`s for redundancy requires iterating over their fields, which is
// problematic if some of those fields point to recursive `Union`s. To avoid cycles,
// compare `TypedDict`s by name/identity instead of using the `has_relation_to`
// machinery.
if let (Type::TypedDict(element_td), Type::TypedDict(ty_td)) = (element_type, ty) {
if element_td == ty_td {
return;
}
continue;
}
if should_simplify_full && !matches!(element_type, Type::TypeAlias(_)) {
if ty.is_redundant_with(self.db, element_type) {
return;

View File

@@ -126,6 +126,16 @@ fn try_metaclass_cycle_initial<'db>(
})
}
fn fields_cycle_initial<'db>(
_db: &'db dyn Db,
_id: salsa::Id,
_self: ClassLiteral<'db>,
_specialization: Option<Specialization<'db>>,
_field_policy: CodeGeneratorKind<'db>,
) -> FxIndexMap<Name, Field<'db>> {
FxIndexMap::default()
}
/// A category of classes with code generation capabilities (with synthesized methods).
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, salsa::Update, get_size2::GetSize)]
pub(crate) enum CodeGeneratorKind<'db> {
@@ -555,8 +565,10 @@ impl<'db> ClassType<'db> {
TypeRelation::Assignability => ConstraintSet::from(!other.is_final(db)),
},
// Protocol and Generic are not represented by a ClassType.
ClassBase::Protocol | ClassBase::Generic => ConstraintSet::from(false),
// Protocol, Generic, and TypedDict are not represented by a ClassType.
ClassBase::Protocol | ClassBase::Generic | ClassBase::TypedDict => {
ConstraintSet::from(false)
}
ClassBase::Class(base) => match (base, other) {
(ClassType::NonGeneric(base), ClassType::NonGeneric(other)) => {
@@ -579,11 +591,6 @@ impl<'db> ClassType<'db> {
ConstraintSet::from(false)
}
},
ClassBase::TypedDict => {
// TODO: Implement subclassing and assignability for TypedDicts.
ConstraintSet::from(true)
}
}
})
}
@@ -2824,7 +2831,10 @@ impl<'db> ClassLiteral<'db> {
/// Returns a list of all annotated attributes defined in this class, or any of its superclasses.
///
/// See [`ClassLiteral::own_fields`] for more details.
#[salsa::tracked(returns(ref), heap_size=get_size2::GetSize::get_heap_size)]
#[salsa::tracked(
returns(ref),
cycle_initial=fields_cycle_initial,
heap_size=get_size2::GetSize::get_heap_size)]
pub(crate) fn fields(
self,
db: &'db dyn Db,
@@ -3933,6 +3943,7 @@ pub enum KnownClass {
SupportsIndex,
Iterable,
Iterator,
Mapping,
// typing_extensions
ExtensionsTypeVar, // must be distinct from typing.TypeVar, backports new features
// Collections
@@ -4047,6 +4058,7 @@ impl KnownClass {
| Self::ABCMeta
| Self::Iterable
| Self::Iterator
| Self::Mapping
// Evaluating `NotImplementedType` in a boolean context was deprecated in Python 3.9
// and raises a `TypeError` in Python >=3.14
// (see https://docs.python.org/3/library/constants.html#NotImplemented)
@@ -4133,6 +4145,7 @@ impl KnownClass {
| KnownClass::SupportsIndex
| KnownClass::Iterable
| KnownClass::Iterator
| KnownClass::Mapping
| KnownClass::ChainMap
| KnownClass::Counter
| KnownClass::DefaultDict
@@ -4218,6 +4231,7 @@ impl KnownClass {
| KnownClass::SupportsIndex
| KnownClass::Iterable
| KnownClass::Iterator
| KnownClass::Mapping
| KnownClass::ChainMap
| KnownClass::Counter
| KnownClass::DefaultDict
@@ -4303,6 +4317,7 @@ impl KnownClass {
| KnownClass::SupportsIndex
| KnownClass::Iterable
| KnownClass::Iterator
| KnownClass::Mapping
| KnownClass::ChainMap
| KnownClass::Counter
| KnownClass::DefaultDict
@@ -4419,7 +4434,8 @@ impl KnownClass {
| Self::BuiltinFunctionType
| Self::ProtocolMeta
| Self::Template
| KnownClass::Path => false,
| Self::Path
| Self::Mapping => false,
}
}
@@ -4492,6 +4508,7 @@ impl KnownClass {
| KnownClass::SupportsIndex
| KnownClass::Iterable
| KnownClass::Iterator
| KnownClass::Mapping
| KnownClass::ChainMap
| KnownClass::Counter
| KnownClass::DefaultDict
@@ -4582,6 +4599,7 @@ impl KnownClass {
Self::Super => "super",
Self::Iterable => "Iterable",
Self::Iterator => "Iterator",
Self::Mapping => "Mapping",
// For example, `typing.List` is defined as `List = _Alias()` in typeshed
Self::StdlibAlias => "_Alias",
// This is the name the type of `sys.version_info` has in typeshed,
@@ -4880,6 +4898,7 @@ impl KnownClass {
| Self::StdlibAlias
| Self::Iterable
| Self::Iterator
| Self::Mapping
| Self::ProtocolMeta
| Self::SupportsIndex => KnownModule::Typing,
Self::TypeAliasType
@@ -5010,6 +5029,7 @@ impl KnownClass {
| Self::InitVar
| Self::Iterable
| Self::Iterator
| Self::Mapping
| Self::NamedTupleFallback
| Self::NamedTupleLike
| Self::ConstraintSet
@@ -5100,6 +5120,7 @@ impl KnownClass {
| Self::InitVar
| Self::Iterable
| Self::Iterator
| Self::Mapping
| Self::NamedTupleFallback
| Self::NamedTupleLike
| Self::ConstraintSet
@@ -5163,6 +5184,7 @@ impl KnownClass {
"TypeVar" => &[Self::TypeVar, Self::ExtensionsTypeVar],
"Iterable" => &[Self::Iterable],
"Iterator" => &[Self::Iterator],
"Mapping" => &[Self::Mapping],
"ParamSpec" => &[Self::ParamSpec],
"ParamSpecArgs" => &[Self::ParamSpecArgs],
"ParamSpecKwargs" => &[Self::ParamSpecKwargs],
@@ -5304,6 +5326,7 @@ impl KnownClass {
| Self::TypeVarTuple
| Self::Iterable
| Self::Iterator
| Self::Mapping
| Self::ProtocolMeta
| Self::NewType => matches!(module, KnownModule::Typing | KnownModule::TypingExtensions),
Self::Deprecated => matches!(module, KnownModule::Warnings | KnownModule::TypingExtensions),

View File

@@ -8043,6 +8043,10 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
// are handled by the default constructor-call logic (we synthesize a `__new__` method for them
// in `ClassType::own_class_member()`).
class.is_known(self.db(), KnownClass::Tuple) && !class.is_generic()
) || CodeGeneratorKind::TypedDict.matches(
self.db(),
class.class_literal(self.db()).0,
class.class_literal(self.db()).1,
);
// temporary special-casing for all subclasses of `enum.Enum`

View File

@@ -12,7 +12,9 @@ use super::diagnostic::{
report_missing_typed_dict_key,
};
use super::{ApplyTypeMappingVisitor, Type, TypeMapping, visitor};
use crate::types::TypeContext;
use crate::types::constraints::ConstraintSet;
use crate::types::generics::InferableTypeVars;
use crate::types::{HasRelationToVisitor, IsDisjointVisitor, TypeContext, TypeRelation};
use crate::{Db, FxIndexMap};
use ordermap::OrderSet;
@@ -76,6 +78,174 @@ impl<'db> TypedDictType<'db> {
),
}
}
// Subtyping between `TypedDict`s follows the algorithm described at:
// https://typing.python.org/en/latest/spec/typeddict.html#subtyping-between-typeddict-types
pub(super) fn has_relation_to_impl(
self,
db: &'db dyn Db,
target: TypedDictType<'db>,
inferable: InferableTypeVars<'_, 'db>,
relation: TypeRelation<'db>,
relation_visitor: &HasRelationToVisitor<'db>,
disjointness_visitor: &IsDisjointVisitor<'db>,
) -> ConstraintSet<'db> {
// First do a quick nominal check that (if it succeeds) means that we can avoid
// materializing the full `TypedDict` schema for either `self` or `target`.
// This should be cheaper in many cases, and also helps us avoid some cycles.
if self
.defining_class
.is_subclass_of(db, target.defining_class)
{
return ConstraintSet::from(true);
}
let self_items = self.items(db);
let target_items = target.items(db);
// Many rules violations short-circuit with "never", but asking whether one field is
// [relation] to/of another can produce more complicated constraints, and we collect those.
let mut constraints = ConstraintSet::from(true);
for (target_item_name, target_item_field) in target_items {
let field_constraints = if target_item_field.is_required() {
// required target fields
let Some(self_item_field) = self_items.get(target_item_name) else {
// Self is missing a required field.
return ConstraintSet::from(false);
};
if !self_item_field.is_required() {
// A required field is not required in self.
return ConstraintSet::from(false);
}
if target_item_field.is_read_only() {
// For `ReadOnly[]` fields in the target, the corresponding fields in
// self need to have the same assignability/subtyping/etc relation
// individually that we're looking for overall between the
// `TypedDict`s.
self_item_field.declared_ty.has_relation_to_impl(
db,
target_item_field.declared_ty,
inferable,
relation,
relation_visitor,
disjointness_visitor,
)
} else {
if self_item_field.is_read_only() {
// A read-only field can't be assigned to a mutable target.
return ConstraintSet::from(false);
}
// For mutable fields in the target, the relation needs to apply both
// ways, or else mutating the target could violate the structural
// invariants of self. For fully-static types, this is "equivalence".
// For gradual types, it depends on the relation, but mutual
// assignability is "consistency".
self_item_field
.declared_ty
.has_relation_to_impl(
db,
target_item_field.declared_ty,
inferable,
relation,
relation_visitor,
disjointness_visitor,
)
.and(db, || {
target_item_field.declared_ty.has_relation_to_impl(
db,
self_item_field.declared_ty,
inferable,
relation,
relation_visitor,
disjointness_visitor,
)
})
}
} else {
// `NotRequired[]` target fields
if target_item_field.is_read_only() {
// As above, for `NotRequired[]` + `ReadOnly[]` fields in the target. It's
// tempting to refactor things and unify some of these calls to
// `has_relation_to_impl`, but this branch will get more complicated when we
// add support for `closed` and `extra_items` (which is why the rules in the
// spec are structured like they are), and following the structure of the spec
// makes it easier to check the logic here.
if let Some(self_item_field) = self_items.get(target_item_name) {
self_item_field.declared_ty.has_relation_to_impl(
db,
target_item_field.declared_ty,
inferable,
relation,
relation_visitor,
disjointness_visitor,
)
} else {
// Self is missing this not-required, read-only item. However, since all
// `TypedDict`s by default are allowed to have "extra items" of any type
// (until we support `closed` and explicit `extra_items`), this key could
// actually turn out to have a value. To make sure this is type-safe, the
// not-required field in the target needs to be assignable from `object`.
// TODO: `closed` and `extra_items` support will go here.
Type::object().when_assignable_to(
db,
target_item_field.declared_ty,
inferable,
)
}
} else {
// As above, for `NotRequired[]` mutable fields in the target. Again the logic
// is largely the same for now, but it will get more complicated with `closed`
// and `extra_items`.
if let Some(self_item_field) = self_items.get(target_item_name) {
if self_item_field.is_read_only() {
// A read-only field can't be assigned to a mutable target.
return ConstraintSet::from(false);
}
if self_item_field.is_required() {
// A required field can't be assigned to a not-required, mutable field
// in the target, because `del` is allowed on the target field.
return ConstraintSet::from(false);
}
// As above, for mutable fields in the target, the relation needs
// to apply both ways.
self_item_field
.declared_ty
.has_relation_to_impl(
db,
target_item_field.declared_ty,
inferable,
relation,
relation_visitor,
disjointness_visitor,
)
.and(db, || {
target_item_field.declared_ty.has_relation_to_impl(
db,
self_item_field.declared_ty,
inferable,
relation,
relation_visitor,
disjointness_visitor,
)
})
} else {
// Self is missing this not-required, mutable field. This isn't ok if self
// has read-only extra items, which all `TypedDict`s effectively do until
// we support `closed` and explicit `extra_items`. See "A subtle
// interaction between two structural assignability rules prevents
// unsoundness" in `typed_dict.md`.
// TODO: `closed` and `extra_items` support will go here.
ConstraintSet::from(false)
}
}
};
constraints.intersect(db, field_constraints);
if constraints.is_never_satisfied(db) {
return constraints;
}
}
constraints
}
}
pub(crate) fn walk_typed_dict_type<'db, V: visitor::TypeVisitor<'db> + ?Sized>(