[ty] Fall back to Divergent for deeply nested specializations

Author: sharkdp

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

crates/ty_python_semantic/src/types.rs

@@ -68,7 +68,7 @@ use crate::types::signatures::{ParameterForm, walk_signature};
 use crate::types::tuple::{TupleSpec, TupleSpecBuilder};
 pub(crate) use crate::types::typed_dict::{TypedDictParams, TypedDictType, walk_typed_dict_type};
 use crate::types::variance::{TypeVarVariance, 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};
@@ -830,6 +830,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))
     }

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,
@@ -1607,10 +1607,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))
             }
         }

crates/ty_python_semantic/src/types/generics.rs

@@ -1290,6 +1290,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.

crates/ty_python_semantic/src/types/visitor.rs

@@ -295,3 +295,91 @@ pub(super) fn any_over_type<'db>(
     visitor.visit_type(db, ty);
     visitor.found_matching_type.get()
 }
+
+/// Returns the 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`.
+pub(super) fn specialization_depth(db: &dyn Db, ty: Type<'_>) -> usize {
+    struct SpecializationDepthVisitor<'db> {
+        seen_types: RefCell<FxIndexSet<NonAtomicType<'db>>>,
+        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) => {
+                    if !self.seen_types.borrow_mut().insert(non_atomic_type) {
+                        return;
+                    }
+
+                    walk_non_atomic_type(db, non_atomic_type, self);
+                    let child_max_depth = self.max_depth.get();
+
+                    let self_depth: usize =
+                        matches!(non_atomic_type, NonAtomicType::GenericAlias(_)).into();
+
+                    self.max_depth.set(
+                        self.max_depth
+                            .get()
+                            .max(child_max_depth.saturating_add(self_depth)),
+                    );
+                }
+            }
+        }
+    }
+
+    let visitor = SpecializationDepthVisitor {
+        seen_types: RefCell::new(FxIndexSet::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
+        );
+
+        // list[list[int]] | list[list[list[int]]]
+        let union_type =
+            UnionType::from_elements(&db, [list_of_list_of_int, list_of_list_of_list_of_int]);
+        assert_eq!(specialization_depth(&db, union_type), 3);
+    }
+}