[mypyc] Support Python 3.12 type alias syntax (PEP 695)

mypy/checker.py

@@ -125,6 +125,7 @@
     TryStmt,
     TupleExpr,
     TypeAlias,
+    TypeAliasStmt,
     TypeInfo,
     TypeVarExpr,
     UnaryExpr,
@@ -5289,6 +5290,9 @@ def remove_capture_conflicts(self, type_map: TypeMap, inferred_types: dict[Var,
                     if node not in inferred_types or not is_subtype(typ, inferred_types[node]):
                         del type_map[expr]
 
+    def visit_type_alias_stmt(self, o: TypeAliasStmt) -> None:
+        self.expr_checker.accept(o.value)
+
     def make_fake_typeinfo(
         self,
         curr_module_fullname: str,

mypy/checkexpr.py

@@ -411,7 +411,9 @@ def analyze_ref_expr(self, e: RefExpr, lvalue: bool = False) -> Type:
             result = self.alias_type_in_runtime_context(
                 node, ctx=e, alias_definition=e.is_alias_rvalue or lvalue
             )
-        elif isinstance(node, (TypeVarExpr, ParamSpecExpr, TypeVarTupleExpr)):
+        elif isinstance(node, TypeVarExpr):
+            return self.named_type("typing.TypeVar")
+        elif isinstance(node, (ParamSpecExpr, TypeVarTupleExpr)):
             result = self.object_type()
         else:
             if isinstance(node, PlaceholderNode):

mypy/fastparse.py

@@ -1783,7 +1783,13 @@ def visit_TypeAlias(self, n: ast_TypeAlias) -> TypeAliasStmt | AssignmentStmt:
         if NEW_GENERIC_SYNTAX in self.options.enable_incomplete_feature:
             type_params = self.translate_type_params(n.type_params)
             value = self.visit(n.value)
-            node = TypeAliasStmt(self.visit_Name(n.name), type_params, value)
+            # Since the value is evaluated lazily, wrap the value inside a lambda.
+            # This helps mypyc.
+            ret = ReturnStmt(value)
+            self.set_line(ret, n.value)
+            value_func = LambdaExpr(body=Block([ret]))
+            self.set_line(value_func, n.value)
+            node = TypeAliasStmt(self.visit_Name(n.name), type_params, value_func)
             return self.set_line(node, n)
         else:
             self.fail(

mypy/nodes.py

@@ -1653,10 +1653,10 @@ class TypeAliasStmt(Statement):
 
     name: NameExpr
     type_args: list[TypeParam]
-    value: Expression  # Will get translated into a type
+    value: LambdaExpr  # Return value will get translated into a type
     invalid_recursive_alias: bool
 
-    def __init__(self, name: NameExpr, type_args: list[TypeParam], value: Expression) -> None:
+    def __init__(self, name: NameExpr, type_args: list[TypeParam], value: LambdaExpr) -> None:
         super().__init__()
         self.name = name
         self.type_args = type_args

mypy/semanal.py

@@ -3765,6 +3765,10 @@ def analyze_alias(
                     last_tvar_name_with_default = tvar_def.name
                 tvar_defs.append(tvar_def)
 
+            if python_3_12_type_alias:
+                with self.allow_unbound_tvars_set():
+                    rvalue.accept(self)
+
             analyzed, depends_on = analyze_type_alias(
                 typ,
                 self,
@@ -5359,7 +5363,7 @@ def visit_type_alias_stmt(self, s: TypeAliasStmt) -> None:
             tag = self.track_incomplete_refs()
             res, alias_tvars, depends_on, qualified_tvars, empty_tuple_index = self.analyze_alias(
                 s.name.name,
-                s.value,
+                s.value.expr(),
                 allow_placeholder=True,
                 declared_type_vars=type_params,
                 all_declared_type_params_names=all_type_params_names,
@@ -5442,6 +5446,7 @@ def visit_type_alias_stmt(self, s: TypeAliasStmt) -> None:
             current_node = existing.node if existing else alias_node
             assert isinstance(current_node, TypeAlias)
             self.disable_invalid_recursive_aliases(s, current_node, s.value)
+            s.name.accept(self)
         finally:
             self.pop_type_args(s.type_args)
 
@@ -5456,7 +5461,11 @@ def visit_name_expr(self, expr: NameExpr) -> None:
 
     def bind_name_expr(self, expr: NameExpr, sym: SymbolTableNode) -> None:
         """Bind name expression to a symbol table node."""
-        if isinstance(sym.node, TypeVarExpr) and self.tvar_scope.get_binding(sym):
+        if (
+            isinstance(sym.node, TypeVarExpr)
+            and self.tvar_scope.get_binding(sym)
+            and not self.allow_unbound_tvars
+        ):
             self.fail(f'"{expr.name}" is a type variable and only valid in type context', expr)
         elif isinstance(sym.node, PlaceholderNode):
             self.process_placeholder(expr.name, "name", expr)

mypyc/irbuild/builder.py

@@ -24,6 +24,9 @@
     ARG_POS,
     GDEF,
     LDEF,
+    PARAM_SPEC_KIND,
+    TYPE_VAR_KIND,
+    TYPE_VAR_TUPLE_KIND,
     ArgKind,
     CallExpr,
     Decorator,
@@ -44,6 +47,7 @@
     TupleExpr,
     TypeAlias,
     TypeInfo,
+    TypeParam,
     UnaryExpr,
     Var,
 )
@@ -1409,3 +1413,45 @@ 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]:
+    """Create objects representing various kinds of Python 3.12 type parameters.
+
+    The "typing_mod" argument is the "_typing" module object. The type objects
+    are looked up from it.
+
+    The returned list has one item for each "type_args" item, in the same order.
+    Each item is either a TypeVar, TypeVarTuple or ParamSpec instance.
+    """
+    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)
+        if type_param.kind != TYPE_VAR_TUPLE_KIND:
+            # To match runtime semantics, pass infer_variance=True
+            tv = builder.py_call(
+                tvt,
+                [builder.load_str(type_param.name), builder.true()],
+                line,
+                arg_kinds=[ARG_POS, ARG_NAMED],
+                arg_names=[None, "infer_variance"],
+            )
+        else:
+            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

mypyc/irbuild/classdef.py

@@ -7,8 +7,6 @@
 from typing import Callable, Final
 
 from mypy.nodes import (
-    PARAM_SPEC_KIND,
-    TYPE_VAR_KIND,
     TYPE_VAR_TUPLE_KIND,
     AssignmentStmt,
     CallExpr,
@@ -57,7 +55,7 @@
     is_optional_type,
     object_rprimitive,
 )
-from mypyc.irbuild.builder import IRBuilder
+from mypyc.irbuild.builder import IRBuilder, create_type_params
 from mypyc.irbuild.function import (
     gen_property_getter_ir,
     gen_property_setter_ir,
@@ -475,35 +473,20 @@ def make_generic_base_class(
 ) -> 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 = []
-    type_var_imported: Value | None = None
-    for type_param in type_args:
-        unpack = False
-        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(mod, "TypeVar", line)
-                type_var_imported = tvt
-        elif type_param.kind == TYPE_VAR_TUPLE_KIND:
-            tvt = builder.py_get_attr(mod, "TypeVarTuple", line)
-            unpack = True
-        else:
-            assert type_param.kind == PARAM_SPEC_KIND
-            tvt = builder.py_get_attr(mod, "ParamSpec", line)
-        tv = builder.py_call(tvt, [builder.load_str(type_param.name)], line)
-        builder.init_type_var(tv, type_param.name, line)
-        if unpack:
+    tvs = create_type_params(builder, mod, type_args, line)
+    args = []
+    for tv, type_param in zip(tvs, type_args):
+        if type_param.kind == TYPE_VAR_TUPLE_KIND:
             # Evaluate *Ts for a TypeVarTuple
             it = builder.call_c(iter_op, [tv], line)
             tv = builder.call_c(next_op, [it], line)
-        tvs.append(tv)
+        args.append(tv)
+
     gent = builder.py_get_attr(mod, "Generic", line)
-    if len(tvs) == 1:
-        arg = tvs[0]
+    if len(args) == 1:
+        arg = args[0]
     else:
-        arg = builder.new_tuple(tvs, line)
+        arg = builder.new_tuple(args, line)
 
     base = builder.call_c(py_get_item_op, [gent, arg], line)
     return base

mypyc/irbuild/statement.py

@@ -12,6 +12,8 @@
 from typing import Callable, Sequence
 
 from mypy.nodes import (
+    ARG_NAMED,
+    ARG_POS,
     AssertStmt,
     AssignmentStmt,
     AwaitExpr,
@@ -37,6 +39,7 @@
     TempNode,
     TryStmt,
     TupleExpr,
+    TypeAliasStmt,
     WhileStmt,
     WithStmt,
     YieldExpr,
@@ -74,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, create_type_params, int_borrow_friendly_op
 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 (
@@ -105,7 +108,9 @@
     coro_op,
     import_from_many_op,
     import_many_op,
+    import_op,
     send_op,
+    set_type_alias_compute_function_op,
     type_op,
     yield_from_except_op,
 )
@@ -1015,3 +1020,30 @@ 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, s: TypeAliasStmt) -> None:
+    line = s.line
+    # Use "_typing" to avoid importing "typing", as the latter can be expensive.
+    # "_typing" includes everything we need here.
+    mod = builder.call_c(import_op, [builder.load_str("_typing")], line)
+    type_params = create_type_params(builder, mod, s.type_args, s.line)
+
+    type_alias_type = builder.py_get_attr(mod, "TypeAliasType", line)
+    args = [builder.load_str(s.name.name), builder.none()]
+    arg_names: list[str | None] = [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(type_alias_type, 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
+    # Python public APIs don't support this, so we use a C primitive.
+    compute_fn = s.value.accept(builder.visitor)
+    builder.builder.primitive_op(set_type_alias_compute_function_op, [alias, compute_fn], line)
+
+    target = builder.get_assignment_target(s.name)
+    builder.assign(target, alias, line)

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
 

mypyc/lib-rt/CPy.h

@@ -768,6 +768,7 @@ PyObject *CPySingledispatch_RegisterFunction(PyObject *singledispatch_func, PyOb
 
 PyObject *CPy_GetAIter(PyObject *obj);
 PyObject *CPy_GetANext(PyObject *aiter);
+void CPy_SetTypeAliasTypeComputeFunction(PyObject *alias, PyObject *compute_value);
 
 #ifdef __cplusplus
 }

mypyc/lib-rt/misc_ops.c

@@ -940,3 +940,34 @@ PyObject *CPy_GetANext(PyObject *aiter)
 error:
     return NULL;
 }
+
+#ifdef CPY_3_12_FEATURES
+
+// Copied from Python 3.12.3, since this struct is internal to CPython. It defines
+// the structure of typing.TypeAliasType objects. We need it since compute_value is
+// not part of the public API, and we need to set it to match Python runtime semantics.
+//
+// IMPORTANT: This needs to be kept in sync with CPython!
+typedef struct {
+    PyObject_HEAD
+    PyObject *name;
+    PyObject *type_params;
+    PyObject *compute_value;
+    PyObject *value;
+    PyObject *module;
+} typealiasobject;
+
+void CPy_SetTypeAliasTypeComputeFunction(PyObject *alias, PyObject *compute_value) {
+    typealiasobject *obj = (typealiasobject *)alias;
+    if (obj->value != NULL) {
+        Py_DECREF(obj->value);
+    }
+    obj->value = NULL;
+    Py_INCREF(compute_value);
+    if (obj->compute_value != NULL) {
+        Py_DECREF(obj->compute_value);
+    }
+    obj->compute_value = compute_value;
+}
+
+#endif

mypyc/primitives/misc_ops.py

@@ -265,3 +265,15 @@
     return_type=c_pyssize_t_rprimitive,
     error_kind=ERR_NEVER,
 )
+
+# Set the lazy value compute function of an TypeAliasType instance (Python 3.12+).
+# This must only be used as part of initializing the object. Any existing value
+# will be cleared.
+set_type_alias_compute_function_op = custom_primitive_op(
+    name="set_type_alias_compute_function",
+    c_function_name="CPy_SetTypeAliasTypeComputeFunction",
+    # (alias object, value compute function)
+    arg_types=[object_rprimitive, object_rprimitive],
+    return_type=void_rtype,
+    error_kind=ERR_NEVER,
+)

mypyc/test-data/fixtures/ir.py

@@ -45,6 +45,7 @@ def __ne__(self, x: object) -> bool: pass
 
 class type:
     def __init__(self, o: object) -> None: ...
+    def __or__(self, o: object) -> Any: ...
     __name__ : str
     __annotations__: Dict[str, Any]