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

@@ -2,12 +2,13 @@
 
 from __future__ import annotations
 
+import enum
 import itertools
 import time
 from collections import defaultdict
 from contextlib import contextmanager
 from typing import Callable, ClassVar, Final, Iterable, Iterator, List, Optional, Sequence, cast
-from typing_extensions import TypeAlias as _TypeAlias, overload
+from typing_extensions import TypeAlias as _TypeAlias, assert_never, overload
 
 import mypy.checker
 import mypy.errorcodes as codes
@@ -272,6 +273,20 @@ class Finished(Exception):
     """Raised if we can terminate overload argument check early (no match)."""
 
 
+@enum.unique
+class UseReverse(enum.Enum):
+    """Used in `visit_op_expr` to enable or disable reverse method checks."""
+
+    DEFAULT = 0
+    ALWAYS = 1
+    NEVER = 2
+
+
+USE_REVERSE_DEFAULT: Final = UseReverse.DEFAULT
+USE_REVERSE_ALWAYS: Final = UseReverse.ALWAYS
+USE_REVERSE_NEVER: Final = UseReverse.NEVER
+
+
 class ExpressionChecker(ExpressionVisitor[Type]):
     """Expression type checker.
 
@@ -3318,7 +3333,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, *, use_reverse: UseReverse = USE_REVERSE_DEFAULT) -> Type:
         """Type check a binary operator expression."""
         if e.analyzed:
             # It's actually a type expression X | Y.
@@ -3366,6 +3381,27 @@ 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.__ror__` (the same from both runtime and typing perspective)
+                # 3. Do not allow `dict.__ror__` to be executed, since this is a special case
+                proper_right_type = get_proper_type(self.accept(e.right))
+                if isinstance(proper_right_type, TypedDictType):
+                    use_reverse = USE_REVERSE_ALWAYS
+            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__`.
+                proper_right_type = get_proper_type(self.accept(e.right))
+                if is_named_instance(proper_right_type, "builtins.dict"):
+                    use_reverse = USE_REVERSE_NEVER
+
         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 +3421,33 @@ 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)
+            if use_reverse is UseReverse.DEFAULT or use_reverse is UseReverse.NEVER:
+                result, method_type = self.check_op(
+                    method,
+                    base_type=left_type,
+                    arg=e.right,
+                    context=e,
+                    allow_reverse=use_reverse is UseReverse.DEFAULT,
+                )
+            elif use_reverse is UseReverse.ALWAYS:
+                result, method_type = self.check_op(
+                    # The reverse operator here gives better error messages:
+                    # Given:
+                    #   d: dict[int, str]
+                    #   t: YourTD
+                    #   d | t
+                    # Without the switch:
+                    # - Unsupported operand types for | ("YourTD" and "dict[int, str]")
+                    # With the switch:
+                    # - Unsupported operand types for | ("dict[int, str]" and "YourTD")
+                    operators.reverse_op_methods[method],
+                    base_type=self.accept(e.right),
+                    arg=e.left,
+                    context=e,
+                    allow_reverse=False,
+                )
+            else:
+                assert_never(use_reverse)
             e.method_type = method_type
             return result
         else:

test-data/unit/check-typeddict.test

@@ -3236,3 +3236,100 @@ 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 typing import Dict
+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]"
+
+d1: Dict[str, int]
+d2: Dict[int, str]
+
+reveal_type(foo1 | d1)  # N: Revealed type is "builtins.dict[builtins.str, builtins.object]"
+foo1 | d2  # E: Unsupported operand types for | ("Foo" and "Dict[int, str]")
+[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
+
+class SubDict(dict): ...
+
+foo | SubDict()
+[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, Any], /) -> Dict[str, object]
+main:12: error: No overload variant of "__ror__" of "dict" matches argument type "Foo"
+main:12: note: Possible overload variants:
+main:12: note:     def __ror__(self, Dict[Any, Any], /) -> Dict[Any, Any]
+main:12: note:     def [T, T2] __ror__(self, Dict[T, T2], /) -> Dict[Union[Any, T], Union[Any, T2]]
+[builtins fixtures/dict.pyi]
+[typing fixtures/typing-typeddict-iror.pyi]
+
+[case testTypedDictWith__ror__method]
+from typing import Dict
+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: 'a'} | foo  # E: Dict entry 0 has incompatible type "int": "str"; expected "str": "Any"
+
+d1: Dict[str, int]
+d2: Dict[int, str]
+
+reveal_type(d1 | foo)  # N: Revealed type is "builtins.dict[builtins.str, builtins.object]"
+d2 | foo  # E: Unsupported operand types for | ("Dict[int, str]" and "Foo")
+
+1 | foo  # E: Unsupported left operand type for | ("int")
+[builtins fixtures/dict.pyi]
+[typing fixtures/typing-typeddict-iror.pyi]
+
+[case testTypedDictWith__ior__method]
+from typing import Dict
+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"
+
+d1: Dict[str, int]
+d2: Dict[int, str]
+
+foo |= d1  # E: Argument 1 to "__ior__" of "TypedDict" has incompatible type "Dict[str, int]"; expected "Foo"
+foo |= d2  # E: Argument 1 to "__ior__" of "TypedDict" has incompatible type "Dict[int, str]"; expected "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, Any]) -> 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]