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?
 

test-data/unit/check-typeddict.test

@@ -3236,3 +3236,56 @@ 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 "typing.Mapping[builtins.str, builtins.object]"
+reveal_type(foo1 | {})  # N: Revealed type is "typing.Mapping[builtins.str, builtins.object]"
+[builtins fixtures/dict.pyi]
+[typing fixtures/typing-typeddict.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, Mapping[str, object], /) -> Mapping[str, object]
+[builtins fixtures/dict.pyi]
+[typing fixtures/typing-typeddict.pyi]
+
+# TODO: add `__ror__` method check, after `__ror__` definition is fixed.
+
+[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.pyi]

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

@@ -71,3 +71,18 @@ class _TypedDict(Mapping[str, object]):
     def pop(self, k: NoReturn, default: T = ...) -> object: ...
     def update(self: T, __m: T) -> None: ...
     def __delitem__(self, k: NoReturn) -> None: ...
+    # It is a bit of a lie:
+    # 1. it is only supported since 3.9 in runtime
+    # 2. `__or__` uses `dict[str, Any]`, not `Mapping[str, object]`
+    @overload
+    def __or__(self, __value: Self) -> Self: ...
+    @overload
+    def __or__(self, __value: Mapping[str, object]) -> Mapping[str, object]: ...
+    # TODO: re-enable after `__ror__` definition is fixed
+    # https://github.com/python/typeshed/issues/10678
+    # @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]