Support TypeVarTuple and ParamSpec

python · JukkaL · Jun 11, 2024 · Jun 6, 2024 · Jun 7, 2024 · Jun 7, 2024
commit 2841eca9ec3ba24011053978e0067bd90beba703
diff --git a/mypyc/irbuild/classdef.py b/mypyc/irbuild/classdef.py
@@ -24,6 +24,9 @@
     TypeInfo,
     TypeParam,
     is_class_var,
+    TYPE_VAR_KIND,
+    TYPE_VAR_TUPLE_KIND,
+    PARAM_SPEC_KIND,
 )
 from mypy.types import ENUM_REMOVED_PROPS, Instance, RawExpressionType, get_proper_type
 from mypyc.common import PROPSET_PREFIX
@@ -64,7 +67,7 @@
 )
 from mypyc.irbuild.util import dataclass_type, get_func_def, is_constant, is_dataclass_decorator
 from mypyc.primitives.dict_ops import dict_new_op, dict_set_item_op
-from mypyc.primitives.generic_ops import py_hasattr_op, py_setattr_op, py_get_item_op
+from mypyc.primitives.generic_ops import py_hasattr_op, py_setattr_op, py_get_item_op, iter_op, next_op
 from mypyc.primitives.misc_ops import (
     dataclass_sleight_of_hand,
     not_implemented_op,
@@ -459,11 +462,25 @@ def allocate_class(builder: IRBuilder, cdef: ClassDef) -> Value:
 
 
 def make_generic_base_class(builder: IRBuilder, type_args: list[TypeParam], line: int) -> Value:
+    """Construct Generic[...] base class object for a new-style generic class (Python 3.12)."""
     mod = builder.call_c(import_op, [builder.load_str("_typing")], line)
     tvs = []
     for type_param in type_args:
-        tvt = builder.py_get_attr(mod, "TypeVar", line)
+        unpack = False
+        if type_param.kind == TYPE_VAR_KIND:
+            name = "TypeVar"
+        elif type_param.kind == TYPE_VAR_TUPLE_KIND:
+            name = "TypeVarTuple"
+            unpack = True
+        else:
+            assert type_param.kind == PARAM_SPEC_KIND
+            name = "ParamSpec"
+        tvt = builder.py_get_attr(mod, name, line)
         tv = builder.py_call(tvt, [builder.load_str(type_param.name)], line)
+        if unpack:
+            # Evaluate *Ts for a TypeVarTuple
+            it = builder.call_c(iter_op, [tv], line)
+            tv = builder.call_c(next_op, [it], line)
         tvs.append(tv)
     gent = builder.py_get_attr(mod, "Generic", line)
     if len(tvs) == 1:

diff --git a/mypyc/irbuild/statement.py b/mypyc/irbuild/statement.py
@@ -41,6 +41,7 @@
     WithStmt,
     YieldExpr,
     YieldFromExpr,
+    TypeAliasStmt,
 )
 from mypyc.ir.ops import (
     NAMESPACE_MODULE,
@@ -1015,3 +1016,7 @@ def transform_await_expr(builder: IRBuilder, o: AwaitExpr) -> Value:
 
 def transform_match_stmt(builder: IRBuilder, m: MatchStmt) -> None:
     m.accept(MatchVisitor(builder, m))
+
+
+def transform_type_alias_stmt(builder: IRBuilder, m: TypeAliasStmt) -> None:
+    assert False
diff --git a/mypyc/irbuild/visitor.py b/mypyc/irbuild/visitor.py
@@ -137,6 +137,7 @@
     transform_raise_stmt,
     transform_return_stmt,
     transform_try_stmt,
+    transform_type_alias_stmt,
     transform_while_stmt,
     transform_with_stmt,
     transform_yield_expr,
@@ -251,7 +252,7 @@ def visit_match_stmt(self, stmt: MatchStmt) -> None:
         transform_match_stmt(self.builder, stmt)
 
     def visit_type_alias_stmt(self, stmt: TypeAliasStmt) -> None:
-        self.bail('The "type" statement is not yet supported by mypyc', stmt.line)
+        transform_type_alias_stmt(self.builder, stmt)
 
     # Expressions
 

diff --git a/mypyc/test-data/run-python312.test b/mypyc/test-data/run-python312.test
@@ -2,6 +2,8 @@
 # flags: --enable-incomplete-feature=NewGenericSyntax
 from typing import Any, TypeAliasType
 
+from testutil import assertRaises
+
 def id[T](x: T) -> T:
     return x
 
@@ -47,4 +49,35 @@ def test_generic_class_via_any() -> None:
     d2 = d_any[int, str](2, 'y')
     assert d2.x == 2
     assert d2.y == 'y'
+
+    with assertRaises(TypeError):
+        c_any[int, str]
+    with assertRaises(TypeError):
+        d_any[int]
+
+class E[*Ts]: pass
+
+def test_type_var_tuple() -> None:
+    e: E[int, str] = E()
+    e_any: Any = E
+    assert isinstance(e_any(), E)
+    assert isinstance(e_any[int](), E)
+    assert isinstance(e_any[int, str](), E)
+
+class F[**P]: pass
+
+def test_param_spec() -> None:
+    f: F[[int, str]] = F()
+    f_any: Any = F
+    assert isinstance(f_any(), F)
+    assert isinstance(f_any[[int, str]](), F)
+
+#class E: pass
+
+#type A = E
+
+#def test_simple_type_alias() -> None:
+#    assert isinstance(A, TypeAliasType)
+#    assert str(A) == "E"
+
 [typing fixtures/typing-full.pyi]