[ty] Generic "manual" PEP 695 type aliases

Author: sharkdp

crates/ty_ide/src/goto_type_definition.rs

@@ -740,8 +740,26 @@ mod tests {
             "#,
         );
 
-        // TODO: This should jump to the definition of `Alias` above.
-        assert_snapshot!(test.goto_type_definition(), @"No type definitions found");
+        assert_snapshot!(test.goto_type_definition(), @r#"
+        info[goto-type-definition]: Type definition
+         --> main.py:4:1
+          |
+        2 | from typing_extensions import TypeAliasType
+        3 |
+        4 | Alias = TypeAliasType("Alias", tuple[int, int])
+          | ^^^^^
+        5 |
+        6 | Alias
+          |
+        info: Source
+         --> main.py:6:1
+          |
+        4 | Alias = TypeAliasType("Alias", tuple[int, int])
+        5 |
+        6 | Alias
+          | ^^^^^
+          |
+        "#);
     }
 
     #[test]

crates/ty_python_semantic/resources/mdtest/pep613_type_aliases.md

@@ -217,8 +217,7 @@ MyList = TypeAliasType("MyList", list[T], type_params=(T,))
 MyAlias5 = Callable[[MyList[T]], int]
 
 def _(c: MyAlias5[int]):
-    # TODO: should be (list[int], /) -> int
-    reveal_type(c)  # revealed: (Unknown, /) -> int
+    reveal_type(c)  # revealed: (list[int], /) -> int
 
 K = TypeVar("K")
 V = TypeVar("V")
@@ -228,14 +227,12 @@ MyDict = TypeAliasType("MyDict", dict[K, V], type_params=(K, V))
 MyAlias6 = Callable[[MyDict[K, V]], int]
 
 def _(c: MyAlias6[str, bytes]):
-    # TODO: should be (dict[str, bytes], /) -> int
-    reveal_type(c)  # revealed: (Unknown, /) -> int
+    reveal_type(c)  # revealed: (dict[str, bytes], /) -> int
 
 ListOrDict: TypeAlias = MyList[T] | dict[str, T]
 
 def _(x: ListOrDict[int]):
-    # TODO: should be list[int] | dict[str, int]
-    reveal_type(x)  # revealed: Unknown | dict[str, int]
+    reveal_type(x)  # revealed: list[int] | dict[str, int]
 
 MyAlias7: TypeAlias = Callable[Concatenate[T, ...], None]
 

crates/ty_python_semantic/resources/mdtest/pep695_type_aliases.md

@@ -223,8 +223,27 @@ T = TypeVar("T")
 IntAndT = TypeAliasType("IntAndT", tuple[int, T], type_params=(T,))
 
 def f(x: IntAndT[str]) -> None:
-    # TODO: This should be `tuple[int, str]`
-    reveal_type(x)  # revealed: Unknown
+    reveal_type(x)  # revealed: tuple[int, str]
+
+U = TypeVar("U", default=str)
+
+ListOrSet = TypeAliasType("ListOrSet", list[U] | set[U], type_params=(U,))
+
+def g(
+    list_or_set_of_int: ListOrSet[int],
+    list_or_set_of_str: ListOrSet,
+) -> None:
+    reveal_type(list_or_set_of_int)  # revealed: list[int] | set[int]
+    reveal_type(list_or_set_of_str)  # revealed: list[str] | set[str]
+
+MyDict = TypeAliasType("MyDict", dict[U, T], type_params=(U, T))
+
+def h(
+    dict_int_str: MyDict[int, str],
+    dict_str_unknown: MyDict,
+) -> None:
+    reveal_type(dict_int_str)  # revealed: dict[int, str]
+    reveal_type(dict_str_unknown)  # revealed: dict[str, Unknown]
 ```
 
 ### Error cases
@@ -241,6 +260,35 @@ def get_name() -> str:
 IntOrStr = TypeAliasType(get_name(), int | str)
 ```
 
+#### Type parameters argument is not a tuple
+
+```py
+from typing_extensions import TypeAliasType, TypeVar
+
+T = TypeVar("T")
+
+# error: [invalid-type-alias-type] "`type_params` argument to `TypeAliasType` must be a tuple"
+IntAndT = TypeAliasType("IntAndT", tuple[int, T], type_params=T)
+
+def _(x: IntAndT[str]) -> None:
+    reveal_type(x)  # revealed: Unknown
+```
+
+#### Invalid type parameters entries
+
+```py
+from typing_extensions import TypeAliasType, TypeVar
+
+T = TypeVar("T")
+
+# TODO: This should be an error
+IntAndT = TypeAliasType("IntAndT", tuple[int, T], type_params=(str,))
+
+# error: [non-subscriptable] "Cannot subscript non-generic type alias"
+def _(x: IntAndT[str]) -> None:
+    reveal_type(x)  # revealed: Unknown
+```
+
 ## Cyclic aliases
 
 ### Self-referential

crates/ty_python_semantic/src/types.rs

@@ -12898,6 +12898,8 @@ pub struct ManualPEP695TypeAliasType<'db> {
     pub name: ast::name::Name,
     pub definition: Option<Definition<'db>>,
     pub value: Type<'db>,
+    generic_context: Option<GenericContext<'db>>,
+    specialization: Option<Specialization<'db>>,
 }
 
 // The Salsa heap is tracked separately.
@@ -12912,12 +12914,50 @@ fn walk_manual_pep_695_type_alias<'db, V: visitor::TypeVisitor<'db> + ?Sized>(
 }
 
 impl<'db> ManualPEP695TypeAliasType<'db> {
+    pub(crate) fn value_type(self, db: &'db dyn Db) -> Type<'db> {
+        self.apply_function_specialization(db, self.value(db))
+    }
+
+    fn apply_function_specialization(self, db: &'db dyn Db, ty: Type<'db>) -> Type<'db> {
+        if let Some(generic_context) = self.generic_context(db) {
+            let specialization = self
+                .specialization(db)
+                .unwrap_or_else(|| generic_context.default_specialization(db, None));
+            ty.apply_specialization(db, specialization)
+        } else {
+            ty
+        }
+    }
+
+    pub(crate) fn apply_specialization(
+        self,
+        db: &'db dyn Db,
+        f: impl FnOnce(GenericContext<'db>) -> Specialization<'db>,
+    ) -> ManualPEP695TypeAliasType<'db> {
+        match self.generic_context(db) {
+            None => self,
+            Some(generic_context) => {
+                let specialization = f(generic_context);
+                ManualPEP695TypeAliasType::new(
+                    db,
+                    self.name(db),
+                    self.definition(db),
+                    self.value(db),
+                    self.generic_context(db),
+                    Some(specialization),
+                )
+            }
+        }
+    }
+
     fn normalized_impl(self, db: &'db dyn Db, visitor: &NormalizedVisitor<'db>) -> Self {
         Self::new(
             db,
             self.name(db),
             self.definition(db),
             self.value(db).normalized_impl(db, visitor),
+            self.generic_context(db),
+            self.specialization(db),
         )
     }
 
@@ -12929,6 +12969,8 @@ impl<'db> ManualPEP695TypeAliasType<'db> {
             self.definition(db),
             self.value(db)
                 .recursive_type_normalized_impl(db, div, true)?,
+            self.generic_context(db),
+            self.specialization(db),
         ))
     }
 }
@@ -12999,7 +13041,7 @@ impl<'db> TypeAliasType<'db> {
     pub fn value_type(self, db: &'db dyn Db) -> Type<'db> {
         match self {
             TypeAliasType::PEP695(type_alias) => type_alias.value_type(db),
-            TypeAliasType::ManualPEP695(type_alias) => type_alias.value(db),
+            TypeAliasType::ManualPEP695(type_alias) => type_alias.value_type(db),
         }
     }
 
@@ -13018,24 +13060,25 @@ impl<'db> TypeAliasType<'db> {
     }
 
     pub(crate) fn generic_context(self, db: &'db dyn Db) -> Option<GenericContext<'db>> {
-        // TODO: Add support for generic non-PEP695 type aliases.
         match self {
             TypeAliasType::PEP695(type_alias) => type_alias.generic_context(db),
-            TypeAliasType::ManualPEP695(_) => None,
+            TypeAliasType::ManualPEP695(type_alias) => type_alias.generic_context(db),
         }
     }
 
     pub(crate) fn specialization(self, db: &'db dyn Db) -> Option<Specialization<'db>> {
         match self {
             TypeAliasType::PEP695(type_alias) => type_alias.specialization(db),
-            TypeAliasType::ManualPEP695(_) => None,
+            TypeAliasType::ManualPEP695(type_alias) => type_alias.specialization(db),
         }
     }
 
     fn apply_function_specialization(self, db: &'db dyn Db, ty: Type<'db>) -> Type<'db> {
         match self {
             TypeAliasType::PEP695(type_alias) => type_alias.apply_function_specialization(db, ty),
-            TypeAliasType::ManualPEP695(_) => ty,
+            TypeAliasType::ManualPEP695(type_alias) => {
+                type_alias.apply_function_specialization(db, ty)
+            }
         }
     }
 
@@ -13048,7 +13091,9 @@ impl<'db> TypeAliasType<'db> {
             TypeAliasType::PEP695(type_alias) => {
                 TypeAliasType::PEP695(type_alias.apply_specialization(db, f))
             }
-            TypeAliasType::ManualPEP695(_) => self,
+            TypeAliasType::ManualPEP695(type_alias) => {
+                TypeAliasType::ManualPEP695(type_alias.apply_specialization(db, f))
+            }
         }
     }
 }

crates/ty_python_semantic/src/types/class.rs

@@ -19,7 +19,7 @@ use crate::semantic_index::{
 use crate::types::bound_super::BoundSuperError;
 use crate::types::constraints::{ConstraintSet, IteratorConstraintsExtension};
 use crate::types::context::InferContext;
-use crate::types::diagnostic::{INVALID_TYPE_ALIAS_TYPE, SUPER_CALL_IN_NAMED_TUPLE_METHOD};
+use crate::types::diagnostic::SUPER_CALL_IN_NAMED_TUPLE_METHOD;
 use crate::types::enums::enum_metadata;
 use crate::types::function::{DataclassTransformerParams, KnownFunction};
 use crate::types::generics::{
@@ -35,9 +35,9 @@ use crate::types::{
     ApplyTypeMappingVisitor, Binding, BoundSuperType, CallableType, CallableTypes, DATACLASS_FLAGS,
     DataclassFlags, DataclassParams, DeprecatedInstance, FindLegacyTypeVarsVisitor,
     HasRelationToVisitor, IsDisjointVisitor, IsEquivalentVisitor, KnownInstanceType,
-    ManualPEP695TypeAliasType, MaterializationKind, NormalizedVisitor, PropertyInstanceType,
-    StringLiteralType, TypeAliasType, TypeContext, TypeMapping, TypeRelation, TypedDictParams,
-    UnionBuilder, VarianceInferable, binding_type, declaration_type, determine_upper_bound,
+    MaterializationKind, NormalizedVisitor, PropertyInstanceType, StringLiteralType, TypeContext,
+    TypeMapping, TypeRelation, TypedDictParams, UnionBuilder, VarianceInferable, binding_type,
+    declaration_type, determine_upper_bound,
 };
 use crate::{
     Db, FxIndexMap, FxIndexSet, FxOrderSet, Program,
@@ -5664,42 +5664,6 @@ impl KnownClass {
                 )));
             }
 
-            KnownClass::TypeAliasType => {
-                let assigned_to = index
-                    .try_expression(ast::ExprRef::from(call_expression))
-                    .and_then(|expr| expr.assigned_to(db));
-
-                let containing_assignment = assigned_to.as_ref().and_then(|assigned_to| {
-                    match assigned_to.node(module).targets.as_slice() {
-                        [ast::Expr::Name(target)] => Some(index.expect_single_definition(target)),
-                        _ => None,
-                    }
-                });
-
-                let [Some(name), Some(value), ..] = overload.parameter_types() else {
-                    return;
-                };
-
-                let Some(name) = name.as_string_literal() else {
-                    if let Some(builder) =
-                        context.report_lint(&INVALID_TYPE_ALIAS_TYPE, call_expression)
-                    {
-                        builder.into_diagnostic(
-                            "The name of a `typing.TypeAlias` must be a string literal",
-                        );
-                    }
-                    return;
-                };
-                overload.set_return_type(Type::KnownInstance(KnownInstanceType::TypeAliasType(
-                    TypeAliasType::ManualPEP695(ManualPEP695TypeAliasType::new(
-                        db,
-                        ast::name::Name::new(name.value(db)),
-                        containing_assignment,
-                        value,
-                    )),
-                )));
-            }
-
             _ => {}
         }
     }