Allow PEP 604 type aliases to be used in runtime context.

mypy/checker.py

@@ -2112,6 +2112,7 @@ def visit_assignment_stmt(self, s: AssignmentStmt) -> None:
             self.fail(message_registry.DEPENDENT_FINAL_IN_CLASS_BODY, s)
 
     def check_type_alias_rvalue(self, s: AssignmentStmt) -> None:
+        assert s.type_alias is not None
         if not (self.is_stub and isinstance(s.rvalue, OpExpr) and s.rvalue.op == '|'):
             # We do this mostly for compatibility with old semantic analyzer.
             # TODO: should we get rid of this?
@@ -2132,6 +2133,7 @@ def accept_items(e: Expression) -> None:
                     self.expr_checker.accept(e)
 
             accept_items(s.rvalue)
+        s.type_alias.rtype = get_proper_type(alias_type)
         self.store_type(s.lvalues[-1], alias_type)
 
     def check_assignment(self, lvalue: Lvalue, rvalue: Expression, infer_lvalue_type: bool = True,

mypy/checkexpr.py

@@ -3242,6 +3242,13 @@ class LongName(Generic[T]): ...
             return type_object_type(tuple_fallback(item).type, self.named_type)
         elif isinstance(item, AnyType):
             return AnyType(TypeOfAny.from_another_any, source_any=item)
+        elif isinstance(item, UnionType):
+            if alias.rtype is not None:
+                return alias.rtype
+            elif self.chk.options.python_version[:2] >= (3, 10):
+                return self.named_type('types.UnionType')
+            else:
+                return self.named_type('builtins.object')
         else:
             if alias_definition:
                 return AnyType(TypeOfAny.special_form)

mypy/fixup.py

@@ -139,6 +139,8 @@ def visit_var(self, v: Var) -> None:
 
     def visit_type_alias(self, a: TypeAlias) -> None:
         a.target.accept(self.type_fixer)
+        if a.rtype is not None:
+            a.rtype.accept(self.type_fixer)
 
 
 class TypeFixer(TypeVisitor[None]):

mypy/nodes.py

@@ -1099,7 +1099,7 @@ class AssignmentStmt(Statement):
     """
 
     __slots__ = ('lvalues', 'rvalue', 'type', 'unanalyzed_type', 'new_syntax',
-                 'is_alias_def', 'is_final_def')
+                 'is_alias_def', 'type_alias', 'is_final_def')
 
     lvalues: List[Lvalue]
     # This is a TempNode if and only if no rvalue (x: t).
@@ -1112,6 +1112,7 @@ class AssignmentStmt(Statement):
     new_syntax: bool
     # Does this assignment define a type alias?
     is_alias_def: bool
+    type_alias: Optional["TypeAlias"]
     # Is this a final definition?
     # Final attributes can't be re-assigned once set, and can't be overridden
     # in a subclass. This flag is not set if an attempted declaration was found to
@@ -1129,6 +1130,7 @@ def __init__(self, lvalues: List[Lvalue], rvalue: Expression,
         self.unanalyzed_type = type
         self.new_syntax = new_syntax
         self.is_alias_def = False
+        self.type_alias = None
         self.is_final_def = False
 
     def accept(self, visitor: StatementVisitor[T]) -> T:
@@ -2971,14 +2973,15 @@ def f(x: B[T]) -> T: ...  # without T, Any would be used here
         within functions that can't be looked up from the symbol table)
     """
     __slots__ = ('target', '_fullname', 'alias_tvars', 'no_args', 'normalized',
-                 'line', 'column', '_is_recursive', 'eager')
+                 'line', 'column', '_is_recursive', 'eager', 'rtype')
 
     def __init__(self, target: 'mypy.types.Type', fullname: str, line: int, column: int,
                  *,
                  alias_tvars: Optional[List[str]] = None,
                  no_args: bool = False,
                  normalized: bool = False,
-                 eager: bool = False) -> None:
+                 eager: bool = False,
+                 rtype: 'Optional[mypy.types.Type]' = None) -> None:
         self._fullname = fullname
         self.target = target
         if alias_tvars is None:
@@ -2990,6 +2993,7 @@ def __init__(self, target: 'mypy.types.Type', fullname: str, line: int, column:
         # it is the cached value.
         self._is_recursive: Optional[bool] = None
         self.eager = eager
+        self.rtype = rtype
         super().__init__(line, column)
 
     @property
@@ -3010,6 +3014,7 @@ def serialize(self) -> JsonDict:
             "normalized": self.normalized,
             "line": self.line,
             "column": self.column,
+            "rtype": self.rtype.serialize() if self.rtype is not None else None
         }
         return data
 
@@ -3026,8 +3031,9 @@ def deserialize(cls, data: JsonDict) -> 'TypeAlias':
         normalized = data['normalized']
         line = data['line']
         column = data['column']
+        rtype = mypy.types.deserialize_type(data['rtype']) if data['rtype'] is not None else None
         return cls(target, fullname, line, column, alias_tvars=alias_tvars,
-                   no_args=no_args, normalized=normalized)
+                   no_args=no_args, normalized=normalized, rtype=rtype)
 
 
 class PlaceholderNode(SymbolNode):

mypy/semanal.py

@@ -2733,6 +2733,7 @@ def check_and_set_up_type_alias(self, s: AssignmentStmt) -> bool:
                                alias_tvars=alias_tvars,
                                no_args=no_args,
                                eager=eager)
+        s.type_alias = alias_node
         if isinstance(s.rvalue, (IndexExpr, CallExpr)):  # CallExpr is for `void = type(None)`
             s.rvalue.analyzed = TypeAliasExpr(alias_node)
             s.rvalue.analyzed.line = s.line

mypyc/test-data/fixtures/ir.py

@@ -1,6 +1,7 @@
 # These builtins stubs are used implicitly in AST to IR generation
 # test cases.
 
+import types
 from typing import (
     TypeVar, Generic, List, Iterator, Iterable, Dict, Optional, Tuple, Any, Set,
     overload, Mapping, Union, Callable, Sequence,

test-data/unit/check-union-or-syntax.test

@@ -196,3 +196,47 @@ def f(x: Union[int, str, None]) -> None:
     else:
         reveal_type(x)  # N: Revealed type is "None"
 [builtins fixtures/isinstance.pyi]
+
+[case testUnionOrSyntaxMetaclassOverride]
+# flags: --python-version 3.10
+import types
+
+class Meta(type):
+    def __or__(self, other) -> types.UnionType: pass
+
+class F(metaclass=Meta): pass
+
+reveal_type(F | F)  # N: Revealed type is "types.UnionType"
+UF = F | F
+reveal_type(UF)  # N: Revealed type is "types.UnionType"
+[builtins fixtures/tuple.pyi]
+
+[case testUnionOrSyntaxMetaclassOverrideType]
+# flags: --python-version 3.10
+class U(type): pass
+
+class Meta(type):
+    def __or__(self, other) -> U: pass # type: ignore
+
+class F(metaclass=Meta): pass
+
+reveal_type(F | F)  # N: Revealed type is "__main__.U"
+UF = F | F
+reveal_type(UF)  # N: Revealed type is "__main__.U"
+
+[case testUnionOrSyntaxMetaclassOverrideUnion]
+# flags: --python-version 3.10
+from typing import TypeVar, Union, Type
+
+T = TypeVar('T')
+U = TypeVar('U')
+
+class Meta(type):
+    def __or__(self: Type[T], other: Type[U]) -> Union[T, U]: pass
+
+class A(metaclass=Meta): pass
+class B(metaclass=Meta): pass
+
+reveal_type(A | B)  # N: Revealed type is "Union[__main__.A, __main__.B*]"
+C = A | B
+reveal_type(C)  # N: Revealed type is "Union[__main__.A, __main__.B*]"

test-data/unit/fixtures/callable.pyi

@@ -1,3 +1,4 @@
+import types
 from typing import Generic, Tuple, TypeVar, Union
 
 T = TypeVar('T')

test-data/unit/lib-stub/types.pyi

@@ -13,3 +13,5 @@ class ModuleType:
 if sys.version_info >= (3, 10):
     class Union:
         def __or__(self, x) -> Union: ...
+
+    class UnionType: ...