special case TypedDict.get

Author: randolf-scholz

mypy/checkmember.py

@@ -68,6 +68,7 @@
     TypedDictType,
     TypeOfAny,
     TypeType,
+    TypeVarId,
     TypeVarLikeType,
     TypeVarTupleType,
     TypeVarType,
@@ -1400,6 +1401,67 @@ def analyze_typeddict_access(
             fallback=mx.chk.named_type("builtins.function"),
             name=name,
         )
+    elif name == "get":
+        # synthesize TypedDict.get() overloads
+        t = TypeVarType(
+            "T",
+            "T",
+            id=TypeVarId(-1),
+            values=[],
+            upper_bound=mx.chk.named_type("builtins.object"),
+            default=AnyType(TypeOfAny.from_omitted_generics),
+        )
+        str_type = mx.chk.named_type("builtins.str")
+        fn_type = mx.chk.named_type("builtins.function")
+        object_type = mx.chk.named_type("builtins.object")
+
+        overloads: list[CallableType] = []
+        # add two overloads per TypedDictType spec
+        for key, val in typ.items.items():
+            # first overload: def(Literal[key]) -> val
+            no_default = CallableType(
+                arg_types=[LiteralType(key, fallback=str_type)],
+                arg_kinds=[ARG_POS],
+                arg_names=[None],
+                ret_type=val,
+                fallback=fn_type,
+                name=name,
+            )
+            # second Overload: def [T] (Literal[key], default: T | Val, /) -> T | Val
+            with_default = CallableType(
+                variables=[t],
+                arg_types=[LiteralType(key, fallback=str_type), UnionType.make_union([val, t])],
+                arg_kinds=[ARG_POS, ARG_POS],
+                arg_names=[None, None],
+                ret_type=UnionType.make_union([val, t]),
+                fallback=fn_type,
+                name=name,
+            )
+            overloads.append(no_default)
+            overloads.append(with_default)
+
+        # finally, add fallback overloads when a key is used that is not in the TypedDict
+        # def (str) -> object
+        fallback_no_default = CallableType(
+            arg_types=[str_type],
+            arg_kinds=[ARG_POS],
+            arg_names=[None],
+            ret_type=object_type,
+            fallback=fn_type,
+            name=name,
+        )
+        # def (str, object) -> object
+        fallback_with_default = CallableType(
+            arg_types=[str_type, object_type],
+            arg_kinds=[ARG_POS, ARG_POS],
+            arg_names=[None, None],
+            ret_type=object_type,
+            fallback=fn_type,
+            name=name,
+        )
+        overloads.append(fallback_no_default)
+        overloads.append(fallback_with_default)
+        return Overloaded(overloads)
     return _analyze_member_access(name, typ.fallback, mx, override_info)
 
 

test-data/unit/check-typeddict.test

@@ -1016,7 +1016,7 @@ class A: pass
 D = TypedDict('D', {'x': List[int], 'y': int})
 d: D
 reveal_type(d.get('x', [])) # N: Revealed type is "builtins.list[builtins.int]"
-d.get('x', ['x']) # E: List item 0 has incompatible type "str"; expected "int"
+reveal_type(d.get('x', ['x']))  # N: Revealed type is "Union[builtins.list[builtins.int], builtins.list[builtins.str]]"
 a = ['']
 reveal_type(d.get('x', a)) # N: Revealed type is "Union[builtins.list[builtins.int], builtins.list[builtins.str]]"
 [builtins fixtures/dict.pyi]
@@ -1026,14 +1026,22 @@ reveal_type(d.get('x', a)) # N: Revealed type is "Union[builtins.list[builtins.i
 from typing import TypedDict
 D = TypedDict('D', {'x': int, 'y': str})
 d: D
-d.get() # E: All overload variants of "get" of "Mapping" require at least one argument \
+d.get() # E: All overload variants of "get" require at least one argument \
         # N: Possible overload variants: \
-        # N:     def get(self, k: str) -> object \
-        # N:     def [V] get(self, k: str, default: object) -> object
-d.get('x', 1, 2) # E: No overload variant of "get" of "Mapping" matches argument types "str", "int", "int" \
+        # N:     def get(Literal['x'], /) -> int \
+        # N:     def [T] get(Literal['x'], Union[int, T], /) -> Union[int, T] \
+        # N:     def get(Literal['y'], /) -> str \
+        # N:     def [T] get(Literal['y'], Union[str, T], /) -> Union[str, T] \
+        # N:     def get(str, /) -> object \
+        # N:     def get(str, object, /) -> object
+d.get('x', 1, 2) # E: No overload variant of "get" matches argument types "str", "int", "int" \
                  # N: Possible overload variants: \
-                 # N:     def get(self, k: str) -> object \
-                 # N:     def [V] get(self, k: str, default: Union[int, V]) -> object
+                 # N:     def get(Literal['x'], /) -> int \
+                 # N:     def [T] get(Literal['x'], Union[int, T], /) -> Union[int, T] \
+                 # N:     def get(Literal['y'], /) -> str \
+                 # N:     def [T] get(Literal['y'], Union[int, T], /) -> Union[str, T] \
+                 # N:     def get(str, /) -> object \
+                 # N:     def get(str, object, /) -> object
 x = d.get('z')
 reveal_type(x) # N: Revealed type is "builtins.object"
 s = ''

test-data/unit/pythoneval.test

@@ -1046,11 +1046,14 @@ reveal_type(d.get(s))
 _testTypedDictGet.py:6: note: Revealed type is "Union[builtins.int, None]"
 _testTypedDictGet.py:7: note: Revealed type is "Union[builtins.str, None]"
 _testTypedDictGet.py:8: note: Revealed type is "builtins.object"
-_testTypedDictGet.py:9: error: All overload variants of "get" of "Mapping" require at least one argument
+_testTypedDictGet.py:9: error: All overload variants of "get" require at least one argument
 _testTypedDictGet.py:9: note: Possible overload variants:
-_testTypedDictGet.py:9: note:     def get(self, str, /) -> object
-_testTypedDictGet.py:9: note:     def get(self, str, /, default: object) -> object
-_testTypedDictGet.py:9: note:     def [_T] get(self, str, /, default: _T) -> object
+_testTypedDictGet.py:9: note:     def get(Literal['x'], /) -> int
+_testTypedDictGet.py:9: note:     def [T] get(Literal['x'], Union[int, T], /) -> Union[int, T]
+_testTypedDictGet.py:9: note:     def get(Literal['y'], /) -> str
+_testTypedDictGet.py:9: note:     def [T] get(Literal['y'], Union[str, T], /) -> Union[str, T]
+_testTypedDictGet.py:9: note:     def get(str, /) -> object
+_testTypedDictGet.py:9: note:     def get(str, object, /) -> object
 _testTypedDictGet.py:11: note: Revealed type is "builtins.object"
 
 [case testTypedDictMappingMethods]