Fix issues with type aliases and new style unions (#14181)

Fix aliases like this and other aliases involving new-style unions:
```
A = type[int] | str
```

Fixes #12392. Fixes #14158.

Author: JukkaL

mypy/checker.py

@@ -2668,26 +2668,7 @@ def visit_assignment_stmt(self, s: AssignmentStmt) -> None:
             self.msg.annotation_in_unchecked_function(context=s)
 
     def check_type_alias_rvalue(self, s: AssignmentStmt) -> 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?
-            alias_type = self.expr_checker.accept(s.rvalue)
-        else:
-            # Avoid type checking 'X | Y' in stubs, since there can be errors
-            # on older Python targets.
-            alias_type = AnyType(TypeOfAny.special_form)
-
-            def accept_items(e: Expression) -> None:
-                if isinstance(e, OpExpr) and e.op == "|":
-                    accept_items(e.left)
-                    accept_items(e.right)
-                else:
-                    # Nested union types have been converted to type context
-                    # in semantic analysis (such as in 'list[int | str]'),
-                    # so we don't need to deal with them here.
-                    self.expr_checker.accept(e)
-
-            accept_items(s.rvalue)
+        alias_type = self.expr_checker.accept(s.rvalue)
         self.store_type(s.lvalues[-1], alias_type)
 
     def check_assignment(

mypy/checkexpr.py

@@ -2847,6 +2847,9 @@ def visit_ellipsis(self, e: EllipsisExpr) -> Type:
 
     def visit_op_expr(self, e: OpExpr) -> Type:
         """Type check a binary operator expression."""
+        if e.analyzed:
+            # It's actually a type expression X | Y.
+            return self.accept(e.analyzed)
         if e.op == "and" or e.op == "or":
             return self.check_boolean_op(e, e)
         if e.op == "*" and isinstance(e.left, ListExpr):

mypy/nodes.py

@@ -1969,10 +1969,20 @@ def accept(self, visitor: ExpressionVisitor[T]) -> T:
 
 
 class OpExpr(Expression):
-    """Binary operation (other than . or [] or comparison operators,
-    which have specific nodes)."""
+    """Binary operation.
 
-    __slots__ = ("op", "left", "right", "method_type", "right_always", "right_unreachable")
+    The dot (.), [] and comparison operators have more specific nodes.
+    """
+
+    __slots__ = (
+        "op",
+        "left",
+        "right",
+        "method_type",
+        "right_always",
+        "right_unreachable",
+        "analyzed",
+    )
 
     __match_args__ = ("left", "op", "right")
 
@@ -1985,15 +1995,20 @@ class OpExpr(Expression):
     right_always: bool
     # Per static analysis only: Is the right side unreachable?
     right_unreachable: bool
+    # Used for expressions that represent a type "X | Y" in some contexts
+    analyzed: TypeAliasExpr | None
 
-    def __init__(self, op: str, left: Expression, right: Expression) -> None:
+    def __init__(
+        self, op: str, left: Expression, right: Expression, analyzed: TypeAliasExpr | None = None
+    ) -> None:
         super().__init__()
         self.op = op
         self.left = left
         self.right = right
         self.method_type = None
         self.right_always = False
         self.right_unreachable = False
+        self.analyzed = analyzed
 
     def accept(self, visitor: ExpressionVisitor[T]) -> T:
         return visitor.visit_op_expr(self)

mypy/semanal.py

@@ -3472,7 +3472,11 @@ def check_and_set_up_type_alias(self, s: AssignmentStmt) -> bool:
             no_args=no_args,
             eager=eager,
         )
-        if isinstance(s.rvalue, (IndexExpr, CallExpr)):  # CallExpr is for `void = type(None)`
+        if isinstance(s.rvalue, (IndexExpr, CallExpr, OpExpr)) and (
+            not isinstance(rvalue, OpExpr)
+            or (self.options.python_version >= (3, 10) or self.is_stub_file)
+        ):
+            # Note: CallExpr is for "void = type(None)" and OpExpr is for "X | Y" union syntax.
             s.rvalue.analyzed = TypeAliasExpr(alias_node)
             s.rvalue.analyzed.line = s.line
             # we use the column from resulting target, to get better location for errors

mypy/server/aststrip.py

@@ -54,6 +54,7 @@
     MypyFile,
     NameExpr,
     Node,
+    OpExpr,
     OverloadedFuncDef,
     RefExpr,
     StarExpr,
@@ -222,6 +223,10 @@ def visit_index_expr(self, node: IndexExpr) -> None:
         node.analyzed = None  # May have been an alias or type application.
         super().visit_index_expr(node)
 
+    def visit_op_expr(self, node: OpExpr) -> None:
+        node.analyzed = None  # May have been an alias
+        super().visit_op_expr(node)
+
     def strip_ref_expr(self, node: RefExpr) -> None:
         node.kind = None
         node.node = None

mypy/strconv.py

@@ -413,6 +413,8 @@ def visit_call_expr(self, o: mypy.nodes.CallExpr) -> str:
         return self.dump(a + extra, o)
 
     def visit_op_expr(self, o: mypy.nodes.OpExpr) -> str:
+        if o.analyzed:
+            return o.analyzed.accept(self)
         return self.dump([o.op, o.left, o.right], o)
 
     def visit_comparison_expr(self, o: mypy.nodes.ComparisonExpr) -> str:

mypy/traverser.py

@@ -262,6 +262,8 @@ def visit_call_expr(self, o: CallExpr) -> None:
     def visit_op_expr(self, o: OpExpr) -> None:
         o.left.accept(self)
         o.right.accept(self)
+        if o.analyzed is not None:
+            o.analyzed.accept(self)
 
     def visit_comparison_expr(self, o: ComparisonExpr) -> None:
         for operand in o.operands:

mypy/treetransform.py

@@ -519,7 +519,12 @@ def visit_call_expr(self, node: CallExpr) -> CallExpr:
         )
 
     def visit_op_expr(self, node: OpExpr) -> OpExpr:
-        new = OpExpr(node.op, self.expr(node.left), self.expr(node.right))
+        new = OpExpr(
+            node.op,
+            self.expr(node.left),
+            self.expr(node.right),
+            cast(Optional[TypeAliasExpr], self.optional_expr(node.analyzed)),
+        )
         new.method_type = self.optional_type(node.method_type)
         return new
 

test-data/unit/check-type-aliases.test

@@ -948,6 +948,17 @@ c.SpecialExplicit = 4
 [builtins fixtures/tuple.pyi]
 [typing fixtures/typing-medium.pyi]
 
+[case testNewStyleUnionInTypeAliasWithMalformedInstance]
+# flags: --python-version 3.10
+from typing import List
+
+A = List[int, str] | int  # E: "list" expects 1 type argument, but 2 given
+B = int | list[int, str]  # E: "list" expects 1 type argument, but 2 given
+a: A
+b: B
+reveal_type(a)  # N: Revealed type is "Union[builtins.list[Any], builtins.int]"
+reveal_type(b)  # N: Revealed type is "Union[builtins.int, builtins.list[Any]]"
+
 [case testValidTypeAliasValues]
 from typing import TypeVar, Generic, List
 

test-data/unit/fine-grained.test

@@ -10277,3 +10277,34 @@ A = str
 m.py:5: error: Invalid statement in TypedDict definition; expected "field_name: field_type"
 ==
 m.py:5: error: Invalid statement in TypedDict definition; expected "field_name: field_type"
+
+[case testTypeAliasWithNewStyleUnionChangedToVariable]
+# flags: --python-version 3.10
+import a
+
+[file a.py]
+from b import C, D
+A = C | D
+a: A
+reveal_type(a)
+
+[file b.py]
+C = int
+D = str
+
+[file b.py.2]
+C = "x"
+D = "y"
+
+[file b.py.3]
+C = str
+D = int
+[out]
+a.py:4: note: Revealed type is "Union[builtins.int, builtins.str]"
+==
+a.py:2: error: Unsupported left operand type for | ("str")
+a.py:3: error: Variable "a.A" is not valid as a type
+a.py:3: note: See https://mypy.readthedocs.io/en/stable/common_issues.html#variables-vs-type-aliases
+a.py:4: note: Revealed type is "A?"
+==
+a.py:4: note: Revealed type is "Union[builtins.str, builtins.int]"

test-data/unit/pythoneval.test

@@ -1663,7 +1663,7 @@ _testNarrowTypeForDictKeys.py:16: note: Revealed type is "Union[builtins.str, No
 
 [case testTypeAliasWithNewStyleUnion]
 # flags: --python-version 3.10
-from typing import Literal, Type, TypeAlias
+from typing import Literal, Type, TypeAlias, TypeVar
 
 Foo = Literal[1, 2]
 reveal_type(Foo)
@@ -1682,15 +1682,44 @@ Opt4 = float | None
 
 A = Type[int] | str
 B: TypeAlias = Type[int] | str
+C = type[int] | str
+
+D = type[int] | str
+x: D
+reveal_type(x)
+E: TypeAlias = type[int] | str
+y: E
+reveal_type(y)
+F = list[type[int] | str]
+
+T = TypeVar("T", int, str)
+def foo(x: T) -> T:
+    A = type[int] | str
+    a: A
+    return x
 [out]
 _testTypeAliasWithNewStyleUnion.py:5: note: Revealed type is "typing._SpecialForm"
+_testTypeAliasWithNewStyleUnion.py:25: note: Revealed type is "Union[Type[builtins.int], builtins.str]"
+_testTypeAliasWithNewStyleUnion.py:28: note: Revealed type is "Union[Type[builtins.int], builtins.str]"
 
 [case testTypeAliasWithNewStyleUnionInStub]
 # flags: --python-version 3.7
 import m
+a: m.A
+reveal_type(a)
+b: m.B
+reveal_type(b)
+c: m.C
+reveal_type(c)
+d: m.D
+reveal_type(d)
+e: m.E
+reveal_type(e)
+f: m.F
+reveal_type(f)
 
 [file m.pyi]
-from typing import Type
+from typing import Type, Callable
 from typing_extensions import Literal, TypeAlias
 
 Foo = Literal[1, 2]
@@ -1710,8 +1739,27 @@ Opt4 = float | None
 
 A = Type[int] | str
 B: TypeAlias = Type[int] | str
+C = type[int] | str
+reveal_type(C)
+D: TypeAlias = type[int] | str
+E = str | type[int]
+F: TypeAlias = str | type[int]
+G = list[type[int] | str]
+H = list[str | type[int]]
+
+CU1 = int | Callable[[], str | bool]
+CU2: TypeAlias = int | Callable[[], str | bool]
+CU3 = int | Callable[[str | bool], str]
+CU4: TypeAlias = int | Callable[[str | bool], str]
 [out]
 m.pyi:5: note: Revealed type is "typing._SpecialForm"
+m.pyi:22: note: Revealed type is "typing._SpecialForm"
+_testTypeAliasWithNewStyleUnionInStub.py:4: note: Revealed type is "Union[Type[builtins.int], builtins.str]"
+_testTypeAliasWithNewStyleUnionInStub.py:6: note: Revealed type is "Union[Type[builtins.int], builtins.str]"
+_testTypeAliasWithNewStyleUnionInStub.py:8: note: Revealed type is "Union[Type[builtins.int], builtins.str]"
+_testTypeAliasWithNewStyleUnionInStub.py:10: note: Revealed type is "Union[Type[builtins.int], builtins.str]"
+_testTypeAliasWithNewStyleUnionInStub.py:12: note: Revealed type is "Union[builtins.str, Type[builtins.int]]"
+_testTypeAliasWithNewStyleUnionInStub.py:14: note: Revealed type is "Union[builtins.str, Type[builtins.int]]"
 
 [case testEnumNameWorkCorrectlyOn311]
 # flags: --python-version 3.11
@@ -1736,6 +1784,23 @@ _testEnumNameWorkCorrectlyOn311.py:13: note: Revealed type is "Literal['X']?"
 _testEnumNameWorkCorrectlyOn311.py:14: note: Revealed type is "builtins.int"
 _testEnumNameWorkCorrectlyOn311.py:15: note: Revealed type is "builtins.int"
 
+[case testTypeAliasNotSupportedWithNewStyleUnion]
+# flags: --python-version 3.9
+from typing_extensions import TypeAlias
+A = type[int] | str
+B = str | type[int]
+C = str | int
+D: TypeAlias = str | int
+[out]
+_testTypeAliasNotSupportedWithNewStyleUnion.py:3: error: Invalid type alias: expression is not a valid type
+_testTypeAliasNotSupportedWithNewStyleUnion.py:3: error: Value of type "Type[type]" is not indexable
+_testTypeAliasNotSupportedWithNewStyleUnion.py:4: error: Invalid type alias: expression is not a valid type
+_testTypeAliasNotSupportedWithNewStyleUnion.py:4: error: Value of type "Type[type]" is not indexable
+_testTypeAliasNotSupportedWithNewStyleUnion.py:5: error: Invalid type alias: expression is not a valid type
+_testTypeAliasNotSupportedWithNewStyleUnion.py:5: error: Unsupported left operand type for | ("Type[str]")
+_testTypeAliasNotSupportedWithNewStyleUnion.py:6: error: Invalid type alias: expression is not a valid type
+_testTypeAliasNotSupportedWithNewStyleUnion.py:6: error: Unsupported left operand type for | ("Type[str]")
+
 [case testTypedDictUnionGetFull]
 from typing import Dict
 from typing_extensions import TypedDict