Handle recursive PEP 613 type aliases

Author: sharkdp

crates/ty_python_semantic/resources/mdtest/attributes.md

@@ -2617,6 +2617,19 @@ class D:
 reveal_type(D().x)  # revealed: Unknown | tuple[Divergent, Literal[1]]
 ```
 
+And it also works for homogeneous tuples:
+
+```py
+def make_homogeneous_tuple(x: T) -> tuple[T, ...]:
+    return (x, x)
+
+class E:
+    def f(self, other: "E"):
+        self.x = make_homogeneous_tuple(other.x)
+
+reveal_type(E().x)  # revealed: Unknown | tuple[Divergent, ...]
+```
+
 ## Attributes of standard library modules that aren't yet defined
 
 For attributes of stdlib modules that exist in future versions, we can give better diagnostics.

crates/ty_python_semantic/resources/mdtest/pep613_type_aliases.md

@@ -0,0 +1,17 @@
+# PEP 613 type aliases
+
+We do not support PEP 613 type aliases yet. For now, just make sure that we don't panic:
+
+```py
+from typing import TypeAlias
+
+RecursiveTuple: TypeAlias = tuple[int | "RecursiveTuple", str]
+
+def _(rec: RecursiveTuple):
+    reveal_type(rec)  # revealed: tuple[Divergent, str]
+
+RecursiveHomogeneousTuple: TypeAlias = tuple[int | "RecursiveHomogeneousTuple", ...]
+
+def _(rec: RecursiveHomogeneousTuple):
+    reveal_type(rec)  # revealed: tuple[Divergent, ...]
+```

crates/ty_python_semantic/src/types/infer/builder/type_expression.rs

@@ -14,7 +14,8 @@ use crate::types::visitor::any_over_type;
 use crate::types::{
     CallableType, DynamicType, IntersectionBuilder, KnownClass, KnownInstanceType,
     LintDiagnosticGuard, Parameter, Parameters, SpecialFormType, SubclassOfType, Type,
-    TypeAliasType, TypeContext, TypeIsType, UnionBuilder, UnionType, todo_type,
+    TypeAliasType, TypeContext, TypeIsType, UnionBuilder, UnionType,
+    exceeds_max_specialization_depth, todo_type,
 };
 
 /// Type expressions
@@ -560,8 +561,16 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
             ast::Expr::Tuple(elements) => {
                 if let [element, ellipsis @ ast::Expr::EllipsisLiteral(_)] = &*elements.elts {
                     self.infer_expression(ellipsis, TypeContext::default());
-                    let result =
-                        TupleType::homogeneous(self.db(), self.infer_type_expression(element));
+                    let element_ty = self.infer_type_expression(element);
+
+                    let result = TupleType::homogeneous(
+                        self.db(),
+                        if exceeds_max_specialization_depth(self.db(), element_ty) {
+                            Type::divergent()
+                        } else {
+                            element_ty
+                        },
+                    );
                     self.store_expression_type(tuple_slice, Type::tuple(Some(result)));
                     return Some(result);
                 }
@@ -584,7 +593,13 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
                             // TODO: emit a diagnostic
                         }
                     } else {
-                        element_types.push(element_ty);
+                        element_types.push(
+                            if exceeds_max_specialization_depth(self.db(), element_ty) {
+                                Type::divergent()
+                            } else {
+                                element_ty
+                            },
+                        );
                     }
                 }
 

crates/ty_python_semantic/src/types/instance.rs

@@ -72,11 +72,11 @@ impl<'db> Type<'db> {
     {
         Type::tuple(TupleType::heterogeneous(
             db,
-            elements.into_iter().map(Into::into).map(|ty| {
-                if exceeds_max_specialization_depth(db, ty) {
+            elements.into_iter().map(Into::into).map(|element| {
+                if exceeds_max_specialization_depth(db, element) {
                     Type::divergent()
                 } else {
-                    ty
+                    element
                 }
             }),
         ))

crates/ty_python_semantic/src/types/visitor.rs

@@ -388,62 +388,62 @@ mod tests {
     fn test_generics_layering_depth() {
         let db = setup_db();
 
-        let list_of_int =
-            KnownClass::List.to_specialized_instance(&db, [KnownClass::Int.to_instance(&db)]);
+        let int = || KnownClass::Int.to_instance(&db);
+        let list = |element| KnownClass::List.to_specialized_instance(&db, [element]);
+        let dict = |key, value| KnownClass::Dict.to_specialized_instance(&db, [key, value]);
+        let set = |element| KnownClass::Set.to_specialized_instance(&db, [element]);
+        let str = || KnownClass::Str.to_instance(&db);
+        let bytes = || KnownClass::Bytes.to_instance(&db);
+
+        let list_of_int = list(int());
         assert_eq!(specialization_depth(&db, list_of_int), 1);
 
-        let list_of_list_of_int = KnownClass::List.to_specialized_instance(&db, [list_of_int]);
+        let list_of_list_of_int = list(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]);
+        let list_of_list_of_list_of_int = list(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(dict(str(), list_of_int))), 3);
+
         assert_eq!(
-            specialization_depth(&db, set_of_dict_of_str_and_list_of_int),
+            specialization_depth(
+                &db,
+                UnionType::from_elements(&db, [list_of_list_of_list_of_int, list_of_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);
+        assert_eq!(
+            specialization_depth(
+                &db,
+                UnionType::from_elements(&db, [list_of_list_of_int, list_of_list_of_list_of_int])
+            ),
+            3
+        );
 
-        let tuple_of_tuple_of_int = Type::heterogeneous_tuple(
-            &db,
-            [Type::heterogeneous_tuple(
+        assert_eq!(
+            specialization_depth(
                 &db,
-                [KnownClass::Int.to_instance(&db)],
-            )],
+                Type::heterogeneous_tuple(&db, [Type::heterogeneous_tuple(&db, [int()])])
+            ),
+            2
         );
-        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);
+        assert_eq!(
+            specialization_depth(&db, Type::heterogeneous_tuple(&db, [list_of_int, str()])),
+            2
+        );
 
-        let list_of_union_of_lists = KnownClass::List.to_specialized_instance(
-            &db,
-            [UnionType::from_elements(
+        assert_eq!(
+            specialization_depth(
                 &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)]),
-                ],
-            )],
+                list(UnionType::from_elements(
+                    &db,
+                    [list(int()), list(str()), list(bytes())]
+                ))
+            ),
+            2
         );
-        assert_eq!(specialization_depth(&db, list_of_union_of_lists), 2);
     }
 }