Fix mypyc crash with enum type aliases

mypyc/irbuild/main.py

@@ -73,7 +73,7 @@ def build_ir(
 
     for module in modules:
         # First pass to determine free symbols.
-        pbv = PreBuildVisitor(errors, module, singledispatch_info.decorators_to_remove)
+        pbv = PreBuildVisitor(errors, module, singledispatch_info.decorators_to_remove, types)
         module.accept(pbv)
 
         # Construct and configure builder objects (cyclic runtime dependency).

mypyc/irbuild/missingtypevisitor.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+from mypy.nodes import Expression, Node
+from mypy.traverser import ExtendedTraverserVisitor
+from mypy.types import AnyType, Type, TypeOfAny
+
+
+class MissingTypesVisitor(ExtendedTraverserVisitor):
+    """AST visitor that can be used to add any missing types as a generic AnyType."""
+
+    def __init__(self, types: dict[Expression, Type]) -> None:
+        super().__init__()
+        self.types: dict[Expression, Type] = types
+
+    def visit(self, o: Node) -> bool:
+        if isinstance(o, Expression) and o not in self.types:
+            self.types[o] = AnyType(TypeOfAny.special_form)
+
+        # If returns True, will continue to nested nodes.
+        return True

mypyc/irbuild/prebuildvisitor.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from mypy.nodes import (
+    AssignmentStmt,
     Block,
     Decorator,
     Expression,
@@ -16,7 +17,9 @@
     Var,
 )
 from mypy.traverser import ExtendedTraverserVisitor
+from mypy.types import Type
 from mypyc.errors import Errors
+from mypyc.irbuild.missingtypevisitor import MissingTypesVisitor
 
 
 class PreBuildVisitor(ExtendedTraverserVisitor):
@@ -39,6 +42,7 @@ def __init__(
         errors: Errors,
         current_file: MypyFile,
         decorators_to_remove: dict[FuncDef, list[int]],
+        types: dict[Expression, Type],
     ) -> None:
         super().__init__()
         # Dict from a function to symbols defined directly in the
@@ -82,11 +86,20 @@ def __init__(
 
         self.current_file: MypyFile = current_file
 
+        self.missing_types_visitor = MissingTypesVisitor(types)
+
     def visit(self, o: Node) -> bool:
         if not isinstance(o, Import):
             self._current_import_group = None
         return True
 
+    def visit_assignment_stmt(self, stmt: AssignmentStmt) -> None:
+        # These are cases where mypy may not have types for certain expressions,
+        # but mypyc needs some form type to exist.
+        if stmt.is_alias_def:
+            stmt.rvalue.accept(self.missing_types_visitor)
+        return super().visit_assignment_stmt(stmt)
+
     def visit_block(self, block: Block) -> None:
         self._current_import_group = None
         super().visit_block(block)

mypyc/test-data/irbuild-classes.test

@@ -1335,3 +1335,13 @@ def outer():
 if True:
     class OtherInner:  # E: Nested class definitions not supported
         pass
+
+[case testEnumClassAlias]
+from enum import Enum
+from typing import Literal, Union
+
+class SomeEnum(Enum):
+    AVALUE = "a"
+
+ALIAS = Literal[SomeEnum.AVALUE]
+ALIAS2 = Union[Literal[SomeEnum.AVALUE], None]

mypyc/test-data/run-python312.test

@@ -1,5 +1,6 @@
 [case testPEP695Basics]
-from typing import Any, TypeAliasType, cast
+from enum import Enum
+from typing import Any, Literal, TypeAliasType, cast
 
 from testutil import assertRaises
 
@@ -188,6 +189,13 @@ type R = int | list[R]
 def test_recursive_type_alias() -> None:
     assert isinstance(R, TypeAliasType)
     assert getattr(R, "__value__") == (int | list[R])
+
+class SomeEnum(Enum):
+    AVALUE = "a"
+
+type EnumLiteralAlias1 = Literal[SomeEnum.AVALUE]
+type EnumLiteralAlias2 = Literal[SomeEnum.AVALUE] | None
+EnumLiteralAlias3 = Literal[SomeEnum.AVALUE] | None
 [typing fixtures/typing-full.pyi]
 
 [case testPEP695GenericTypeAlias]