Add |= and | operators support for TypedDict

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,21 @@ def visit_op_expr(self, e: OpExpr) -> Type:
                         return proper_left_type.copy_modified(
                             items=proper_left_type.items + [UnpackType(mapped)]
                         )
+        if is_named_instance(proper_left_type, "builtins.dict") and e.op == "|":
+            # This is a special case for `dict | TypedDict`.
+            # Before this change this operation was not allowed ude to typing limitations,
+            # however, it does perfect sense from runtime's point of view.
+            # So, what 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 TYPE_VAR_TUPLE in self.chk.options.enable_incomplete_feature:
             # Handle tuple[X, ...] + tuple[Y, Z] = tuple[*tuple[X, ...], Y, Z].
             if (
@@ -3385,7 +3400,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

@@ -3249,10 +3249,10 @@ 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]"
+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.pyi]
+[typing fixtures/typing-typeddict-iror.pyi]
 
 [case testTypedDictWith__or__method_error]
 from mypy_extensions import TypedDict
@@ -3267,11 +3267,25 @@ foo | 1
 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]
+main:8: note:     def __or__(self, Dict[str, object], /) -> Dict[str, object]
 [builtins fixtures/dict.pyi]
-[typing fixtures/typing-typeddict.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}
 
-# TODO: add `__ror__` method check, after `__ror__` definition is fixed.
+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
@@ -3288,4 +3302,4 @@ foo |= {'key': 'a', 'b': 'a'}  # E: Extra key "b" for TypedDict "Foo" \
 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]
+[typing fixtures/typing-typeddict-iror.pyi]

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]