Support generic type aliases

python · JukkaL · Jun 17, 2024 · Jun 12, 2024 · Jun 14, 2024 · Jun 14, 2024
commit fecfa80f45272482e75a9382dc598c4ec241c10d
diff --git a/mypyc/irbuild/builder.py b/mypyc/irbuild/builder.py
@@ -46,6 +46,10 @@
     TypeInfo,
     UnaryExpr,
     Var,
+    TypeParam,
+    TYPE_VAR_KIND,
+    TYPE_VAR_TUPLE_KIND,
+    PARAM_SPEC_KIND,
 )
 from mypy.types import (
     AnyType,
@@ -1409,3 +1413,26 @@ def get_call_target_fullname(ref: RefExpr) -> str:
         if isinstance(target, Instance):
             return target.type.fullname
     return ref.fullname
+
+
+def create_type_params(builder: IRBuilder, typing_mod: Value,
+                       type_args: list[TypeParam], line: int) -> list[Value]:
+    tvs = []
+    type_var_imported: Value | None = None
+    for type_param in type_args:
+        if type_param.kind == TYPE_VAR_KIND:
+            if type_var_imported:
+                # Reuse previously imported value as a minor optimization
+                tvt = type_var_imported
+            else:
+                tvt = builder.py_get_attr(typing_mod, "TypeVar", line)
+                type_var_imported = tvt
+        elif type_param.kind == TYPE_VAR_TUPLE_KIND:
+            tvt = builder.py_get_attr(typing_mod, "TypeVarTuple", line)
+        else:
+            assert type_param.kind == PARAM_SPEC_KIND
+            tvt = builder.py_get_attr(typing_mod, "ParamSpec", line)
+        tv = builder.py_call(tvt, [builder.load_str(type_param.name)], line)
+        builder.init_type_var(tv, type_param.name, line)
+        tvs.append(tv)
+    return tvs
diff --git a/mypyc/irbuild/statement.py b/mypyc/irbuild/statement.py
@@ -12,6 +12,8 @@
 from typing import Callable, Sequence
 
 from mypy.nodes import (
+    ARG_POS,
+    ARG_NAMED,
     AssertStmt,
     AssignmentStmt,
     AwaitExpr,
@@ -75,7 +77,7 @@
     object_rprimitive,
 )
 from mypyc.irbuild.ast_helpers import is_borrow_friendly_expr, process_conditional
-from mypyc.irbuild.builder import IRBuilder, int_borrow_friendly_op
+from mypyc.irbuild.builder import IRBuilder, int_borrow_friendly_op, create_type_params
 from mypyc.irbuild.for_helpers import for_loop_helper
 from mypyc.irbuild.generator import add_raise_exception_blocks_to_generator_class
 from mypyc.irbuild.nonlocalcontrol import (
@@ -1024,12 +1026,21 @@ def transform_type_alias_stmt(builder: IRBuilder, s: TypeAliasStmt) -> None:
     line = s.line
     # Use _typing.TypeAliasType to avoid importing "typing", as this can be expensive.
     mod = builder.call_c(import_op, [builder.load_str("_typing")], line)
+    type_params = create_type_params(builder, mod, s.type_args, s.line)
     typ = builder.py_get_attr(mod, "TypeAliasType", line)
 
+    args = [builder.load_str(s.name.name), builder.none()]
+    arg_names = [None, None]
+    arg_kinds = [ARG_POS, ARG_POS]
+    if s.type_args:
+        args.append(builder.new_tuple(type_params, line))
+        arg_names.append("type_params")
+        arg_kinds.append(ARG_NAMED)
+    alias = builder.py_call(typ, args, line, arg_names=arg_names, arg_kinds=arg_kinds)
+
     # Use primitive to set function used to lazily compute type alias type value.
     # The value needs to be lazily computed to match Python runtime behavior, but
     # the public API doesn't support this, so we use a C primitive.
-    alias = builder.py_call(typ, [builder.load_str(s.name.name), builder.none()], line)
     compute_fn = s.value.accept(builder.visitor)
     builder.builder.primitive_op(set_type_alias_compute_function_op, [alias, compute_fn], line)
 

diff --git a/mypyc/test-data/fixtures/typing-full.pyi b/mypyc/test-data/fixtures/typing-full.pyi
@@ -10,6 +10,9 @@ from abc import abstractmethod, ABCMeta
 
 class GenericMeta(type): pass
 
+class _SpecialForm:
+    def __getitem__(self, index): ...
+
 cast = 0
 overload = 0
 Any = 0
@@ -19,7 +22,6 @@ TypeVar = 0
 Generic = 0
 Protocol = 0
 Tuple = 0
-Callable = 0
 _promote = 0
 NamedTuple = 0
 Type = 0
@@ -30,6 +32,7 @@ Literal = 0
 TypedDict = 0
 NoReturn = 0
 NewType = 0
+Callable: _SpecialForm
 
 T = TypeVar('T')
 T_co = TypeVar('T_co', covariant=True)

diff --git a/mypyc/test-data/run-python312.test b/mypyc/test-data/run-python312.test
@@ -190,3 +190,36 @@ def test_recursive_type_alias() -> None:
     assert isinstance(R, TypeAliasType)
     assert getattr(R, "__value__") == (int | list[R])
 [typing fixtures/typing-full.pyi]
+
+[case testPEP695GenericTypeAlias]
+# flags: --enable-incomplete-feature=NewGenericSyntax
+from typing import Callable
+from types import GenericAlias
+
+from testutil import assertRaises
+
+type A[T] = list[T]
+
+def test_generic_alias() -> None:
+    assert type(A[str]) is GenericAlias
+    assert str(A[str]) == "A[str]"
+    assert str(getattr(A, "__value__")) == "list[~T]"
+
+type B[T, S] = dict[S, T]
+
+def test_generic_alias_with_two_args() -> None:
+    assert str(B[str, int]) == "B[str, int]"
+    assert str(getattr(B, "__value__")) == "dict[~S, ~T]"  # TODO: T, S
+
+type C[*Ts] = tuple[*Ts]
+
+def test_type_var_tuple_type_alias() -> None:
+    assert str(C[int, str]) == "C[int, str]"
+    assert str(getattr(C, "__value__")) == "tuple[typing.Unpack[Ts]]"
+
+type D[**P] = Callable[P, int]
+
+def test_param_spec_type_alias() -> None:
+    assert str(D[int, str]) == "D[int, str]"
+    assert str(getattr(D, "__value__")) == "typing.Callable[~P, int]"
+[typing fixtures/typing-full.pyi]
diff --git a/test-data/unit/lib-stub/types.pyi b/test-data/unit/lib-stub/types.pyi
@@ -15,3 +15,5 @@ if sys.version_info >= (3, 10):
 
     class NoneType:
         ...
+
+class GenericAlias: ...
Original file line number	Diff line number	Diff line change
Expand Up		@@ -15,3 +15,5 @@ if sys.version_info >= (3, 10):

		class NoneType:
		...

		class GenericAlias: ...