Fix mypyc crash with enum type aliases (#18725)

mypyc was crashing because it couldn't find the type in the type map.
This PR adds a generic AnyType to the type map if an expression isn't in
the map already.

Tried actually changing mypy to accept these type alias expressions, but
ran into problems with nested type aliases where the inner one doesn't
have the "analyzed" value and ending up with wrong results.

fixes mypyc/mypyc#1064

Author: svalentin

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]