diff --git a/mypyc/irbuild/main.py b/mypyc/irbuild/main.py index 15928d939cbf..7cdc6b686778 100644 --- a/mypyc/irbuild/main.py +++ b/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). diff --git a/mypyc/irbuild/missingtypevisitor.py b/mypyc/irbuild/missingtypevisitor.py new file mode 100644 index 000000000000..e655d270a4a4 --- /dev/null +++ b/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 diff --git a/mypyc/irbuild/prebuildvisitor.py b/mypyc/irbuild/prebuildvisitor.py index 5f178a290138..e630fed0d85a 100644 --- a/mypyc/irbuild/prebuildvisitor.py +++ b/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) diff --git a/mypyc/test-data/irbuild-classes.test b/mypyc/test-data/irbuild-classes.test index 605ab46181e2..ed7c167d8621 100644 --- a/mypyc/test-data/irbuild-classes.test +++ b/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] diff --git a/mypyc/test-data/run-python312.test b/mypyc/test-data/run-python312.test index a5a3f058d1e2..5c0a807c375a 100644 --- a/mypyc/test-data/run-python312.test +++ b/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]