[ty] Fall back to `Divergent` for deeply nested specializations

This commit is contained in:
David Peter 2025-10-18 18:59:36 +02:00
parent e1cada1ec3
commit d2f73a404a
5 changed files with 213 additions and 4 deletions

View File

@ -2457,6 +2457,31 @@ class Counter:
reveal_type(Counter().count) # revealed: Unknown | int
```
We also handle infinitely nested generics:
```py
class NestedLists:
def __init__(self: "NestedLists"):
self.x = 1
def f(self: "NestedLists"):
self.x = [self.x]
reveal_type(NestedLists().x) # revealed: Unknown | Literal[1] | list[Divergent]
class NestedMixed:
def f(self: "NestedMixed"):
self.x = [self.x]
def g(self: "NestedMixed"):
self.x = {self.x}
def h(self: "NestedMixed"):
self.x = {"a": self.x}
reveal_type(NestedMixed().x) # revealed: Unknown | list[Divergent] | set[Divergent] | dict[Unknown | str, Divergent]
```
### Builtin types attributes
This test can probably be removed eventually, but we currently include it because we do not yet

View File

@ -69,7 +69,7 @@ use crate::types::tuple::{TupleSpec, TupleSpecBuilder};
pub(crate) use crate::types::typed_dict::{TypedDictParams, TypedDictType, walk_typed_dict_type};
pub use crate::types::variance::TypeVarVariance;
use crate::types::variance::VarianceInferable;
use crate::types::visitor::any_over_type;
use crate::types::visitor::{any_over_type, specialization_depth};
use crate::unpack::EvaluationMode;
use crate::{Db, FxOrderSet, Module, Program};
pub(crate) use class::{ClassLiteral, ClassType, GenericAlias, KnownClass};
@ -831,6 +831,10 @@ impl<'db> Type<'db> {
Self::Dynamic(DynamicType::Divergent(DivergentType { scope }))
}
pub(crate) const fn is_divergent(&self) -> bool {
matches!(self, Type::Dynamic(DynamicType::Divergent(_)))
}
pub const fn is_unknown(&self) -> bool {
matches!(self, Type::Dynamic(DynamicType::Unknown))
}

View File

@ -37,7 +37,7 @@ use crate::types::{
IsDisjointVisitor, IsEquivalentVisitor, KnownInstanceType, ManualPEP695TypeAliasType,
MaterializationKind, NormalizedVisitor, PropertyInstanceType, StringLiteralType, TypeAliasType,
TypeContext, TypeMapping, TypeRelation, TypedDictParams, UnionBuilder, VarianceInferable,
declaration_type, determine_upper_bound, infer_definition_types,
declaration_type, determine_upper_bound, infer_definition_types, specialization_depth,
};
use crate::{
Db, FxIndexMap, FxIndexSet, FxOrderSet, Program,
@ -1609,10 +1609,34 @@ impl<'db> ClassLiteral<'db> {
db: &'db dyn Db,
f: impl FnOnce(GenericContext<'db>) -> Specialization<'db>,
) -> ClassType<'db> {
// To prevent infinite recursion during type inference for infinite types, we fall back to
// `C[Divergent]` once a certain amount of levels of specialization have occurred. For
// example:
//
// ```py
// x = 1
// while random_bool():
// x = [x]
//
// reveal_type(x) # Unknown | Literal[1] | list[Divergent]
// ```
const MAX_SPECIALIZATION_DEPTH: usize = 10;
match self.generic_context(db) {
None => ClassType::NonGeneric(self),
Some(generic_context) => {
let specialization = f(generic_context);
let mut specialization = f(generic_context);
for (idx, ty) in specialization.types(db).iter().enumerate() {
if specialization_depth(db, *ty) > MAX_SPECIALIZATION_DEPTH {
specialization = specialization.with_replaced_type(
db,
idx,
Type::divergent(self.body_scope(db)),
);
}
}
ClassType::Generic(GenericAlias::new(db, self, specialization))
}
}

View File

@ -1264,6 +1264,27 @@ impl<'db> Specialization<'db> {
// A tuple's specialization will include all of its element types, so we don't need to also
// look in `self.tuple`.
}
/// Returns a copy of this specialization with the type at a given index replaced.
pub(crate) fn with_replaced_type(
self,
db: &'db dyn Db,
index: usize,
new_type: Type<'db>,
) -> Self {
debug_assert!(index < self.types(db).len());
let mut new_types: Box<[_]> = self.types(db).to_vec().into_boxed_slice();
new_types[index] = new_type;
Self::new(
db,
self.generic_context(db),
new_types,
self.materialization_kind(db),
self.tuple_inner(db),
)
}
}
/// A mapping between type variables and types.

View File

@ -1,3 +1,5 @@
use rustc_hash::FxHashMap;
use crate::{
Db, FxIndexSet,
types::{
@ -16,7 +18,10 @@ use crate::{
walk_typed_dict_type, walk_typeis_type, walk_union,
},
};
use std::cell::{Cell, RefCell};
use std::{
cell::{Cell, RefCell},
collections::hash_map::Entry,
};
/// A visitor trait that recurses into nested types.
///
@ -295,3 +300,133 @@ pub(super) fn any_over_type<'db>(
visitor.visit_type(db, ty);
visitor.found_matching_type.get()
}
/// Returns the maximum number of layers of generic specializations for a given type.
///
/// For example, `int` has a depth of `0`, `list[int]` has a depth of `1`, and `list[set[int]]`
/// has a depth of `2`. A set-theoretic type like `list[int] | list[list[int]]` has a maximum
/// depth of `2`.
pub(super) fn specialization_depth(db: &dyn Db, ty: Type<'_>) -> usize {
struct SpecializationDepthVisitor<'db> {
seen_types: RefCell<FxHashMap<NonAtomicType<'db>, Option<usize>>>,
max_depth: Cell<usize>,
}
impl<'db> TypeVisitor<'db> for SpecializationDepthVisitor<'db> {
fn should_visit_lazy_type_attributes(&self) -> bool {
false
}
fn visit_type(&self, db: &'db dyn Db, ty: Type<'db>) {
match TypeKind::from(ty) {
TypeKind::Atomic => {
if ty.is_divergent() {
self.max_depth.set(usize::MAX);
}
}
TypeKind::NonAtomic(non_atomic_type) => {
match self.seen_types.borrow_mut().entry(non_atomic_type) {
Entry::Occupied(cached_depth) => {
self.max_depth
.update(|current| current.max(cached_depth.get().unwrap_or(0)));
return;
}
Entry::Vacant(entry) => {
entry.insert(None);
}
}
let self_depth: usize =
matches!(non_atomic_type, NonAtomicType::GenericAlias(_)).into();
let previous_max_depth = self.max_depth.replace(0);
walk_non_atomic_type(db, non_atomic_type, self);
self.max_depth.update(|max_child_depth| {
previous_max_depth.max(max_child_depth.saturating_add(self_depth))
});
self.seen_types
.borrow_mut()
.insert(non_atomic_type, Some(self.max_depth.get()));
}
}
}
}
let visitor = SpecializationDepthVisitor {
seen_types: RefCell::new(FxHashMap::default()),
max_depth: Cell::new(0),
};
visitor.visit_type(db, ty);
visitor.max_depth.get()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{db::tests::setup_db, types::KnownClass};
#[test]
fn test_generics_layering_depth() {
let db = setup_db();
let list_of_int =
KnownClass::List.to_specialized_instance(&db, [KnownClass::Int.to_instance(&db)]);
assert_eq!(specialization_depth(&db, list_of_int), 1);
let list_of_list_of_int = KnownClass::List.to_specialized_instance(&db, [list_of_int]);
assert_eq!(specialization_depth(&db, list_of_list_of_int), 2);
let list_of_list_of_list_of_int =
KnownClass::List.to_specialized_instance(&db, [list_of_list_of_int]);
assert_eq!(specialization_depth(&db, list_of_list_of_list_of_int), 3);
let set_of_dict_of_str_and_list_of_int = KnownClass::Set.to_specialized_instance(
&db,
[KnownClass::Dict
.to_specialized_instance(&db, [KnownClass::Str.to_instance(&db), list_of_int])],
);
assert_eq!(
specialization_depth(&db, set_of_dict_of_str_and_list_of_int),
3
);
let union_type_1 =
UnionType::from_elements(&db, [list_of_list_of_list_of_int, list_of_list_of_int]);
assert_eq!(specialization_depth(&db, union_type_1), 3);
let union_type_2 =
UnionType::from_elements(&db, [list_of_list_of_int, list_of_list_of_list_of_int]);
assert_eq!(specialization_depth(&db, union_type_2), 3);
let tuple_of_tuple_of_int = Type::heterogeneous_tuple(
&db,
[Type::heterogeneous_tuple(
&db,
[KnownClass::Int.to_instance(&db)],
)],
);
assert_eq!(specialization_depth(&db, tuple_of_tuple_of_int), 2);
let tuple_of_list_of_int_and_str = KnownClass::Tuple
.to_specialized_instance(&db, [list_of_int, KnownClass::Str.to_instance(&db)]);
assert_eq!(specialization_depth(&db, tuple_of_list_of_int_and_str), 1);
let list_of_union_of_lists = KnownClass::List.to_specialized_instance(
&db,
[UnionType::from_elements(
&db,
[
KnownClass::List
.to_specialized_instance(&db, [KnownClass::Int.to_instance(&db)]),
KnownClass::List
.to_specialized_instance(&db, [KnownClass::Str.to_instance(&db)]),
KnownClass::List
.to_specialized_instance(&db, [KnownClass::Bytes.to_instance(&db)]),
],
)],
);
assert_eq!(specialization_depth(&db, list_of_union_of_lists), 2);
}
}