diff --git a/crates/ty_python_semantic/resources/mdtest/attributes.md b/crates/ty_python_semantic/resources/mdtest/attributes.md index 94edff19e6..1e4305d5ba 100644 --- a/crates/ty_python_semantic/resources/mdtest/attributes.md +++ b/crates/ty_python_semantic/resources/mdtest/attributes.md @@ -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 diff --git a/crates/ty_python_semantic/src/types.rs b/crates/ty_python_semantic/src/types.rs index d9a51ce5b9..583926254e 100644 --- a/crates/ty_python_semantic/src/types.rs +++ b/crates/ty_python_semantic/src/types.rs @@ -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)) } diff --git a/crates/ty_python_semantic/src/types/class.rs b/crates/ty_python_semantic/src/types/class.rs index 56df9329ed..6221202308 100644 --- a/crates/ty_python_semantic/src/types/class.rs +++ b/crates/ty_python_semantic/src/types/class.rs @@ -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)) } } diff --git a/crates/ty_python_semantic/src/types/generics.rs b/crates/ty_python_semantic/src/types/generics.rs index f401e0df3f..0b2d61426e 100644 --- a/crates/ty_python_semantic/src/types/generics.rs +++ b/crates/ty_python_semantic/src/types/generics.rs @@ -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. diff --git a/crates/ty_python_semantic/src/types/visitor.rs b/crates/ty_python_semantic/src/types/visitor.rs index 51b77432a4..d49ba3ac15 100644 --- a/crates/ty_python_semantic/src/types/visitor.rs +++ b/crates/ty_python_semantic/src/types/visitor.rs @@ -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, Option>>, + max_depth: Cell, + } + + 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); + } +}