Add |= and | operators support for TypedDict

mypy/checker.py

@@ -7452,14 +7452,25 @@ def infer_operator_assignment_method(typ: Type, operator: str) -> tuple[bool, st
     """
     typ = get_proper_type(typ)
     method = operators.op_methods[operator]
+    existing_method = None
     if isinstance(typ, Instance):
-        if operator in operators.ops_with_inplace_method:
-            inplace_method = "__i" + method[2:]
-            if typ.type.has_readable_member(inplace_method):
-                return True, inplace_method
+        existing_method = _find_inplace_method(typ, method, operator)
+    elif isinstance(typ, TypedDictType):
+        existing_method = _find_inplace_method(typ.fallback, method, operator)
+
+    if existing_method is not None:
+        return True, existing_method
     return False, method
 
 
+def _find_inplace_method(inst: Instance, method: str, operator: str) -> str | None:
+    if operator in operators.ops_with_inplace_method:
+        inplace_method = "__i" + method[2:]
+        if inst.type.has_readable_member(inplace_method):
+            return inplace_method
+    return None
+
+
 def is_valid_inferred_type(typ: Type, is_lvalue_final: bool = False) -> bool:
     """Is an inferred type valid and needs no further refinement?
 

mypy/checkexpr.py

@@ -3318,7 +3318,7 @@ def visit_ellipsis(self, e: EllipsisExpr) -> Type:
         """Type check '...'."""
         return self.named_type("builtins.ellipsis")
 
-    def visit_op_expr(self, e: OpExpr) -> Type:
+    def visit_op_expr(self, e: OpExpr, *, allow_reverse: bool = True) -> Type:
         """Type check a binary operator expression."""
         if e.analyzed:
             # It's actually a type expression X | Y.
@@ -3366,6 +3366,29 @@ def visit_op_expr(self, e: OpExpr) -> Type:
                         return proper_left_type.copy_modified(
                             items=proper_left_type.items + [UnpackType(mapped)]
                         )
+
+        if e.op == "|":
+            if is_named_instance(proper_left_type, "builtins.dict"):
+                # This is a special case for `dict | TypedDict`.
+                # Before this change this operation was not allowed due to typing limitations,
+                # however, it makes perfect sense from the runtime's point of view.
+                # So, what do we do now?
+                # 1. Find `dict | TypedDict` case
+                # 2. Switch `dict.__or__` to `TypedDict.__or__` (the same from typing's perspective)
+                # 3. Do not allow `dict.__ror__` to be executed, since this is a special case
+                # This can later be removed if `typeshed` can do this without special casing.
+                # https://github.com/python/mypy/pull/16249
+                proper_right_type = get_proper_type(self.accept(e.right))
+                if isinstance(proper_right_type, TypedDictType):
+                    reverse_op = OpExpr(e.op, e.right, e.left)
+                    reverse_op.set_line(e)
+                    return self.visit_op_expr(reverse_op, allow_reverse=False)
+            if isinstance(proper_left_type, TypedDictType):
+                # This is the reverse case: `TypedDict | dict`,
+                # simply do not allow the reverse checking:
+                # do not call `__dict__.__ror__`.
+                allow_reverse = False
+
         if TYPE_VAR_TUPLE in self.chk.options.enable_incomplete_feature:
             # Handle tuple[X, ...] + tuple[Y, Z] = tuple[*tuple[X, ...], Y, Z].
             if (
@@ -3385,7 +3408,9 @@ def visit_op_expr(self, e: OpExpr) -> Type:
 
         if e.op in operators.op_methods:
             method = operators.op_methods[e.op]
-            result, method_type = self.check_op(method, left_type, e.right, e, allow_reverse=True)
+            result, method_type = self.check_op(
+                method, left_type, e.right, e, allow_reverse=allow_reverse
+            )
             e.method_type = method_type
             return result
         else:

test-data/unit/check-typeddict.test

@@ -3236,3 +3236,70 @@ def foo(x: int) -> Foo: ...
 f: Foo = {**foo("no")}  # E: Argument 1 to "foo" has incompatible type "str"; expected "int"
 [builtins fixtures/dict.pyi]
 [typing fixtures/typing-typeddict.pyi]
+
+
+[case testTypedDictWith__or__method]
+from mypy_extensions import TypedDict
+
+class Foo(TypedDict):
+    key: int
+
+foo1: Foo = {'key': 1}
+foo2: Foo = {'key': 2}
+
+reveal_type(foo1 | foo2)  # N: Revealed type is "TypedDict('__main__.Foo', {'key': builtins.int})"
+reveal_type(foo1 | {'key': 1})  # N: Revealed type is "TypedDict('__main__.Foo', {'key': builtins.int})"
+reveal_type(foo1 | {'key': 'a'})  # N: Revealed type is "builtins.dict[builtins.str, builtins.object]"
+reveal_type(foo1 | {})  # N: Revealed type is "builtins.dict[builtins.str, builtins.object]"
+[builtins fixtures/dict.pyi]
+[typing fixtures/typing-typeddict-iror.pyi]
+
+[case testTypedDictWith__or__method_error]
+from mypy_extensions import TypedDict
+
+class Foo(TypedDict):
+    key: int
+
+foo: Foo = {'key': 1}
+
+foo | 1
+[out]
+main:8: error: No overload variant of "__or__" of "TypedDict" matches argument type "int"
+main:8: note: Possible overload variants:
+main:8: note:     def __or__(self, Foo, /) -> Foo
+main:8: note:     def __or__(self, Dict[str, object], /) -> Dict[str, object]
+[builtins fixtures/dict.pyi]
+[typing fixtures/typing-typeddict-iror.pyi]
+
+[case testTypedDictWith__ror__method]
+from mypy_extensions import TypedDict
+
+class Foo(TypedDict):
+    key: int
+
+foo: Foo = {'key': 1}
+
+reveal_type({'key': 1} | foo)  # N: Revealed type is "TypedDict('__main__.Foo', {'key': builtins.int})"
+reveal_type({'key': 'a'} | foo)  # N: Revealed type is "builtins.dict[builtins.str, builtins.object]"
+reveal_type({} | foo)  # N: Revealed type is "builtins.dict[builtins.str, builtins.object]"
+
+1 | foo  # E: Unsupported left operand type for | ("int")
+[builtins fixtures/dict.pyi]
+[typing fixtures/typing-typeddict-iror.pyi]
+
+[case testTypedDictWith__ior__method]
+from mypy_extensions import TypedDict
+
+class Foo(TypedDict):
+    key: int
+
+foo: Foo = {'key': 1}
+foo |= {'key': 2}
+
+foo |= {}  # E: Missing key "key" for TypedDict "Foo"
+foo |= {'key': 'a', 'b': 'a'}  # E: Extra key "b" for TypedDict "Foo" \
+                               # E: Incompatible types (expression has type "str", TypedDict item "key" has type "int")
+foo |= {'b': 2}  # E: Missing key "key" for TypedDict "Foo" \
+                 # E: Extra key "b" for TypedDict "Foo"
+[builtins fixtures/dict.pyi]
+[typing fixtures/typing-typeddict-iror.pyi]

test-data/unit/fixtures/dict.pyi

@@ -3,10 +3,12 @@
 from _typeshed import SupportsKeysAndGetItem
 import _typeshed
 from typing import (
-    TypeVar, Generic, Iterable, Iterator, Mapping, Tuple, overload, Optional, Union, Sequence
+    TypeVar, Generic, Iterable, Iterator, Mapping, Tuple, overload, Optional, Union, Sequence,
+    Self,
 )
 
 T = TypeVar('T')
+T2 = TypeVar('T2')
 KT = TypeVar('KT')
 VT = TypeVar('VT')
 
@@ -34,6 +36,21 @@ class dict(Mapping[KT, VT]):
     def get(self, k: KT, default: Union[VT, T]) -> Union[VT, T]: pass
     def __len__(self) -> int: ...
 
+    # This was actually added in 3.9:
+    @overload
+    def __or__(self, __value: dict[KT, VT]) -> dict[KT, VT]: ...
+    @overload
+    def __or__(self, __value: dict[T, T2]) -> dict[Union[KT, T], Union[VT, T2]]: ...
+    @overload
+    def __ror__(self, __value: dict[KT, VT]) -> dict[KT, VT]: ...
+    @overload
+    def __ror__(self, __value: dict[T, T2]) -> dict[Union[KT, T], Union[VT, T2]]: ...
+    # dict.__ior__ should be kept roughly in line with MutableMapping.update()
+    @overload  # type: ignore[misc]
+    def __ior__(self, __value: _typeshed.SupportsKeysAndGetItem[KT, VT]) -> Self: ...
+    @overload
+    def __ior__(self, __value: Iterable[Tuple[KT, VT]]) -> Self: ...
+
 class int: # for convenience
     def __add__(self, x: Union[int, complex]) -> int: pass
     def __radd__(self, x: int) -> int: pass

test-data/unit/fixtures/typing-async.pyi

@@ -24,6 +24,7 @@ ClassVar = 0
 Final = 0
 Literal = 0
 NoReturn = 0
+Self = 0
 
 T = TypeVar('T')
 T_co = TypeVar('T_co', covariant=True)

test-data/unit/fixtures/typing-full.pyi

@@ -30,6 +30,7 @@ Literal = 0
 TypedDict = 0
 NoReturn = 0
 NewType = 0
+Self = 0
 
 T = TypeVar('T')
 T_co = TypeVar('T_co', covariant=True)

test-data/unit/fixtures/typing-medium.pyi

@@ -28,6 +28,7 @@ NoReturn = 0
 NewType = 0
 TypeAlias = 0
 LiteralString = 0
+Self = 0
 
 T = TypeVar('T')
 T_co = TypeVar('T_co', covariant=True)

test-data/unit/fixtures/typing-typeddict-iror.pyi

@@ -0,0 +1,66 @@
+# Test stub for typing module that includes TypedDict `|` operator.
+# It only covers `__or__`, `__ror__`, and `__ior__`.
+#
+# We cannot define these methods in `typing-typeddict.pyi`,
+# because they need `dict` with two type args,
+# and not all tests using `[typing typing-typeddict.pyi]` have the proper
+# `dict` stub.
+#
+# Keep in sync with `typeshed`'s definition.
+from abc import ABCMeta
+
+cast = 0
+assert_type = 0
+overload = 0
+Any = 0
+Union = 0
+Optional = 0
+TypeVar = 0
+Generic = 0
+Protocol = 0
+Tuple = 0
+Callable = 0
+NamedTuple = 0
+Final = 0
+Literal = 0
+TypedDict = 0
+NoReturn = 0
+Required = 0
+NotRequired = 0
+Self = 0
+
+T = TypeVar('T')
+T_co = TypeVar('T_co', covariant=True)
+V = TypeVar('V')
+
+# Note: definitions below are different from typeshed, variances are declared
+# to silence the protocol variance checks. Maybe it is better to use type: ignore?
+
+class Sized(Protocol):
+    def __len__(self) -> int: pass
+
+class Iterable(Protocol[T_co]):
+    def __iter__(self) -> 'Iterator[T_co]': pass
+
+class Iterator(Iterable[T_co], Protocol):
+    def __next__(self) -> T_co: pass
+
+class Sequence(Iterable[T_co]):
+    # misc is for explicit Any.
+    def __getitem__(self, n: Any) -> T_co: pass # type: ignore[misc]
+
+class Mapping(Iterable[T], Generic[T, T_co], metaclass=ABCMeta):
+    pass
+
+# Fallback type for all typed dicts (does not exist at runtime).
+class _TypedDict(Mapping[str, object]):
+    @overload
+    def __or__(self, __value: Self) -> Self: ...
+    @overload
+    def __or__(self, __value: dict[str, object]) -> dict[str, object]: ...
+    @overload
+    def __ror__(self, __value: Self) -> Self: ...
+    @overload
+    def __ror__(self, __value: dict[str, Any]) -> dict[str, object]: ...
+    # supposedly incompatible definitions of __or__ and __ior__
+    def __ior__(self, __value: Self) -> Self: ...  # type: ignore[misc]