fix UnionType.__or__ and add UnionType.__getitem__ (#14687)

Co-authored-by: Alex Waygood <[email protected]>

Author: KotlinIsland

stdlib/@tests/test_cases/check_types.py

@@ -3,9 +3,11 @@
 import sys
 import types
 from collections import UserDict
-from typing import Union
+from typing import Any, Literal, TypeVar, Union
 from typing_extensions import assert_type
 
+_T = TypeVar("_T")
+
 # test `types.SimpleNamespace`
 
 # Valid:
@@ -58,3 +60,13 @@ def foo(self, value: int) -> None:
     @foo.deleter
     def foo(self) -> None:
         self._value = None
+
+
+if sys.version_info > (3, 10):
+    union_type = int | list[_T]
+
+    # ideally this would be `_SpecialForm` (Union)
+    assert_type(union_type | Literal[1], types.UnionType | Any)
+    # Both mypy and pyright special-case this operation,
+    # but in different ways, so we just check that no error is emitted:
+    _ = union_type[int]

stdlib/types.pyi

@@ -717,10 +717,19 @@ if sys.version_info >= (3, 10):
         def __args__(self) -> tuple[Any, ...]: ...
         @property
         def __parameters__(self) -> tuple[Any, ...]: ...
-        def __or__(self, value: Any, /) -> UnionType: ...
-        def __ror__(self, value: Any, /) -> UnionType: ...
+        # `(int | str) | Literal["foo"]` returns a generic alias to an instance of `_SpecialForm` (`Union`).
+        # Normally we'd express this using the return type of `_SpecialForm.__ror__`,
+        # but because `UnionType.__or__` accepts `Any`, type checkers will use
+        # the return type of `UnionType.__or__` to infer the result of this operation
+        # rather than `_SpecialForm.__ror__`. To mitigate this, we use `| Any`
+        # in the return type of `UnionType.__(r)or__`.
+        def __or__(self, value: Any, /) -> UnionType | Any: ...
+        def __ror__(self, value: Any, /) -> UnionType | Any: ...
         def __eq__(self, value: object, /) -> bool: ...
         def __hash__(self) -> int: ...
+        # you can only subscript a `UnionType` instance if at least one of the elements
+        # in the union is a generic alias instance that has a non-empty `__parameters__`
+        def __getitem__(self, parameters: Any) -> object: ...
 
 if sys.version_info >= (3, 13):
     @final