Unwrap type[Union[...]] when solving typevar constraints

Author: sterliakov

mypy/constraints.py

@@ -2,7 +2,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Final, Iterable, List, Sequence
+from typing import TYPE_CHECKING, Final, Iterable, List, Sequence, cast
+from typing_extensions import TypeGuard
 
 import mypy.subtypes
 import mypy.typeops
@@ -339,6 +340,26 @@ def _infer_constraints(
     if isinstance(actual, AnyType) and actual.type_of_any == TypeOfAny.suggestion_engine:
         return []
 
+    # type[A | B] is always represented as type[A] | type[B] internally.
+    # This makes our constraint solver choke on type[T] <: type[A] | type[B],
+    # solving T as generic meet(A, B) which is often `object`. Force unwrap such unions
+    # if both sides are type[...] or unions thereof. See `testTypeVarType` test
+    def _is_type_type(tp: ProperType) -> TypeGuard[TypeType | UnionType]:
+        return (
+            isinstance(tp, TypeType)
+            or isinstance(tp, UnionType)
+            and all(isinstance(get_proper_type(o), TypeType) for o in tp.items)
+        )
+
+    def _unwrap_type_type(tp: TypeType | UnionType) -> ProperType:
+        if isinstance(tp, TypeType):
+            return tp.item
+        return UnionType.make_union([cast(TypeType, o).item for o in tp.items])
+
+    if _is_type_type(template) and _is_type_type(actual):
+        template = _unwrap_type_type(template)
+        actual = _unwrap_type_type(actual)
+
     # If the template is simply a type variable, emit a Constraint directly.
     # We need to handle this case before handling Unions for two reasons:
     #  1. "T <: Union[U1, U2]" is not equivalent to "T <: U1 or T <: U2",

test-data/unit/check-typevar-unbound.test

@@ -69,3 +69,50 @@ from typing import TypeVar
 T = TypeVar("T")
 def f(t: T) -> None:
     a, *b = t  # E: "object" object is not iterable
+
+[case testTypeVarType]
+from typing import Mapping, Type, TypeVar, Union
+T = TypeVar("T")
+
+class A: ...
+class B: ...
+
+lookup_table: Mapping[str, Type[Union[A,B]]]
+def load(lookup_table: Mapping[str, Type[T]], lookup_key: str) -> T:
+    ...
+reveal_type(load(lookup_table, "a"))  # N: Revealed type is "Union[__main__.A, __main__.B]"
+
+lookup_table_a: Mapping[str, Type[A]]
+def load2(lookup_table: Mapping[str, Type[Union[T, int]]], lookup_key: str) -> T:
+    ...
+reveal_type(load2(lookup_table_a, "a"))  # N: Revealed type is "__main__.A"
+
+[builtins fixtures/tuple.pyi]
+
+[case testTypeVarTypeAssignment]
+# Adapted from https://github.com/python/mypy/issues/12115
+from typing import TypeVar, Type, Callable, Union, Any
+
+t1: Type[bool] = bool
+t2: Union[Type[bool], Type[str]] = bool
+
+T1 = TypeVar("T1", bound=Union[bool, str])
+def foo1(t: Type[T1]) -> None: ...
+foo1(t1)
+foo1(t2)
+
+T2 = TypeVar("T2", bool, str)
+def foo2(t: Type[T2]) -> None: ...
+foo2(t1)
+# Rejected correctly: T2 cannot be Union[bool, str]
+foo2(t2)  # E: Value of type variable "T2" of "foo2" cannot be "Union[bool, str]"
+
+T3 = TypeVar("T3")
+def foo3(t: Type[T3]) -> None: ...
+foo3(t1)
+foo3(t2)
+
+def foo4(t: Type[Union[bool, str]]) -> None: ...
+foo4(t1)
+foo4(t2)
+[builtins fixtures/tuple.pyi]