improved TypedDict.get inference

randolf-scholz · randolf-scholz · commit 27e50e25e349 · 2025-09-21T14:00:41.000+02:00
diff --git a/mypy/checkmember.py b/mypy/checkmember.py
@@ -17,6 +17,7 @@
 from mypy.meet import is_overlapping_types
 from mypy.messages import MessageBuilder
 from mypy.nodes import (
+    ARG_OPT,
     ARG_POS,
     ARG_STAR,
     ARG_STAR2,
@@ -68,6 +69,7 @@
     TypedDictType,
     TypeOfAny,
     TypeType,
+    TypeVarId,
     TypeVarLikeType,
     TypeVarTupleType,
     TypeVarType,
@@ -1402,6 +1404,87 @@ def analyze_typeddict_access(
             fallback=mx.chk.named_type("builtins.function"),
             name=name,
         )
+    elif name == "get":
+        # synthesize TypedDict.get() overloads
+        str_type = mx.chk.named_type("builtins.str")
+        fn_type = mx.chk.named_type("builtins.function")
+        object_type = mx.chk.named_type("builtins.object")
+
+        # type variable for default value
+        tvar = TypeVarType(
+            "T",
+            "T",
+            id=TypeVarId(-1),
+            values=[],
+            upper_bound=object_type,
+            default=AnyType(TypeOfAny.from_omitted_generics),
+        )
+        # generate the overloads
+        overloads: list[CallableType] = []
+        for key, value_type in typ.items.items():
+            key_type = LiteralType(key, fallback=str_type)
+
+            if key in typ.required_keys:
+                # If the key is required, we know it must be present in the TypedDict.
+                overload = CallableType(
+                    arg_types=[key_type, object_type],
+                    arg_kinds=[ARG_POS, ARG_OPT],
+                    arg_names=[None, None],
+                    ret_type=value_type,
+                    fallback=fn_type,
+                    name=name,
+                )
+                overloads.append(overload)
+            else:
+                # The key is not required, but if it is present, we know its type.
+                # def (K) -> V | None
+                overload = CallableType(
+                    arg_types=[key_type],
+                    arg_kinds=[ARG_POS],
+                    arg_names=[None],
+                    ret_type=UnionType.make_union([value_type, NoneType()]),
+                    fallback=fn_type,
+                    name=name,
+                )
+                overloads.append(overload)
+
+                # We add an extra overload for the case when the given default is a subtype of the value type.
+                # This makes sure that the return type is inferred as the value type instead of a union.
+                # def (K, V) -> V
+                overload = CallableType(
+                    arg_types=[key_type, value_type],
+                    arg_kinds=[ARG_POS, ARG_POS],
+                    arg_names=[None, None],
+                    ret_type=value_type,
+                    fallback=fn_type,
+                    name=name,
+                )
+                overloads.append(overload)
+
+                # fallback: def [T](K, T) -> V | T
+                overload = CallableType(
+                    variables=[tvar],
+                    arg_types=[key_type, tvar],
+                    arg_kinds=[ARG_POS, ARG_POS],
+                    arg_names=[None, None],
+                    ret_type=UnionType.make_union([value_type, tvar]),
+                    fallback=fn_type,
+                    name=name,
+                )
+                overloads.append(overload)
+
+        # finally, add fallback overload when a key is used that is not in the TypedDict
+        # def (str, object=...) -> object
+        fallback_overload = CallableType(
+            arg_types=[str_type, object_type],
+            arg_kinds=[ARG_POS, ARG_OPT],
+            arg_names=[None, None],
+            ret_type=object_type,
+            fallback=fn_type,
+            name=name,
+        )
+        overloads.append(fallback_overload)
+        return Overloaded(overloads)
     return _analyze_member_access(name, typ.fallback, mx, override_info)
 
 
diff --git a/mypy/plugins/default.py b/mypy/plugins/default.py
@@ -5,7 +5,7 @@
 
 import mypy.errorcodes as codes
 from mypy import message_registry
-from mypy.nodes import DictExpr, IntExpr, StrExpr, UnaryExpr
+from mypy.nodes import IntExpr, StrExpr, UnaryExpr
 from mypy.plugin import (
     AttributeContext,
     ClassDefContext,
@@ -75,6 +75,7 @@
     TypedDictType,
     TypeOfAny,
     TypeVarType,
+    UninhabitedType,
     UnionType,
     get_proper_type,
     get_proper_types,
@@ -120,9 +121,9 @@ def get_function_signature_hook(
     def get_method_signature_hook(
         self, fullname: str
     ) -> Callable[[MethodSigContext], FunctionLike] | None:
-        if fullname == "typing.Mapping.get":
-            return typed_dict_get_signature_callback
-        elif fullname in TD_SETDEFAULT_NAMES:
+        # NOTE: signatures for `__setitem__`, `__delitem__` and `get` are
+        #  defined in checkmember.py/analyze_typeddict_access
+        if fullname in TD_SETDEFAULT_NAMES:
             return typed_dict_setdefault_signature_callback
         elif fullname in TD_POP_NAMES:
             return typed_dict_pop_signature_callback
@@ -212,46 +213,6 @@ def get_class_decorator_hook_2(
         return None
 
 
-def typed_dict_get_signature_callback(ctx: MethodSigContext) -> CallableType:
-    """Try to infer a better signature type for TypedDict.get.
-
-    This is used to get better type context for the second argument that
-    depends on a TypedDict value type.
-    """
-    signature = ctx.default_signature
-    if (
-        isinstance(ctx.type, TypedDictType)
-        and len(ctx.args) == 2
-        and len(ctx.args[0]) == 1
-        and isinstance(ctx.args[0][0], StrExpr)
-        and len(signature.arg_types) == 2
-        and len(signature.variables) == 1
-        and len(ctx.args[1]) == 1
-    ):
-        key = ctx.args[0][0].value
-        value_type = get_proper_type(ctx.type.items.get(key))
-        ret_type = signature.ret_type
-        if value_type:
-            default_arg = ctx.args[1][0]
-            if (
-                isinstance(value_type, TypedDictType)
-                and isinstance(default_arg, DictExpr)
-                and len(default_arg.items) == 0
-            ):
-                # Caller has empty dict {} as default for typed dict.
-                value_type = value_type.copy_modified(required_keys=set())
-            # Tweak the signature to include the value type as context. It's
-            # only needed for type inference since there's a union with a type
-            # variable that accepts everything.
-            tv = signature.variables[0]
-            assert isinstance(tv, TypeVarType)
-            return signature.copy_modified(
-                arg_types=[signature.arg_types[0], make_simplified_union([value_type, tv])],
-                ret_type=ret_type,
-            )
-    return signature
-
-
 def typed_dict_get_callback(ctx: MethodContext) -> Type:
     """Infer a precise return type for TypedDict.get with literal first argument."""
     if (
@@ -263,30 +224,41 @@ def typed_dict_get_callback(ctx: MethodContext) -> Type:
         if keys is None:
             return ctx.default_return_type
 
+        default_type: Type
+        if len(ctx.arg_types) <= 1 or not ctx.arg_types[1]:
+            default_type = NoneType()
+        elif len(ctx.arg_types[1]) == 1 and len(ctx.args[1]) == 1:
+            default_type = ctx.arg_types[1][0]
+        else:
+            return ctx.default_return_type
+
         output_types: list[Type] = []
         for key in keys:
-            value_type = get_proper_type(ctx.type.items.get(key))
+            value_type: Type | None = ctx.type.items.get(key)
             if value_type is None:
                 return ctx.default_return_type
 
-            if len(ctx.arg_types) == 1:
+            if key in ctx.type.required_keys:
                 output_types.append(value_type)
-            elif len(ctx.arg_types) == 2 and len(ctx.arg_types[1]) == 1 and len(ctx.args[1]) == 1:
-                default_arg = ctx.args[1][0]
+            else:
+                # HACK to deal with get(key, {})
+                proper_default = get_proper_type(default_type)
                 if (
-                    isinstance(default_arg, DictExpr)
-                    and len(default_arg.items) == 0
-                    and isinstance(value_type, TypedDictType)
+                    isinstance(vt := get_proper_type(value_type), TypedDictType)
+                    and isinstance(proper_default, Instance)
+                    and proper_default.type.fullname == "builtins.dict"
+                    and len(proper_default.args) == 2
+                    and isinstance(get_proper_type(proper_default.args[0]), UninhabitedType)
+                    and isinstance(get_proper_type(proper_default.args[1]), UninhabitedType)
                 ):
-                    # Special case '{}' as the default for a typed dict type.
-                    output_types.append(value_type.copy_modified(required_keys=set()))
+                    output_types.append(vt.copy_modified(required_keys=set()))
                 else:
                     output_types.append(value_type)
-                    output_types.append(ctx.arg_types[1][0])
-
-        if len(ctx.arg_types) == 1:
-            output_types.append(NoneType())
+                    output_types.append(default_type)
 
+        # for nicer reveal_type, put default at the end, if it is present
+        if default_type in output_types:
+            output_types = [t for t in output_types if t != default_type] + [default_type]
         return make_simplified_union(output_types)
     return ctx.default_return_type
 
diff --git a/test-data/unit/check-literal.test b/test-data/unit/check-literal.test
@@ -1884,7 +1884,7 @@ reveal_type(d[a_key])         # N: Revealed type is "builtins.int"
 reveal_type(d[b_key])         # N: Revealed type is "builtins.str"
 d[c_key]                      # E: TypedDict "Outer" has no key "c"
 
-reveal_type(d.get(a_key, u))  # N: Revealed type is "Union[builtins.int, __main__.Unrelated]"
+reveal_type(d.get(a_key, u))  # N: Revealed type is "builtins.int"
 reveal_type(d.get(b_key, u))  # N: Revealed type is "Union[builtins.str, __main__.Unrelated]"
 reveal_type(d.get(c_key, u))  # N: Revealed type is "builtins.object"
 
@@ -1928,7 +1928,7 @@ u: Unrelated
 reveal_type(a[int_key_good])         # N: Revealed type is "builtins.int"
 reveal_type(b[int_key_good])         # N: Revealed type is "builtins.int"
 reveal_type(c[str_key_good])         # N: Revealed type is "builtins.int"
-reveal_type(c.get(str_key_good, u))  # N: Revealed type is "Union[builtins.int, __main__.Unrelated]"
+reveal_type(c.get(str_key_good, u))  # N: Revealed type is "builtins.int"
 reveal_type(c.get(str_key_bad, u))   # N: Revealed type is "builtins.object"
 
 a[int_key_bad]                       # E: Tuple index out of range
@@ -1993,8 +1993,8 @@ optional_keys: Literal["d", "e"]
 bad_keys: Literal["a", "bad"]
 
 reveal_type(test[good_keys])                      # N: Revealed type is "Union[__main__.A, __main__.B]"
-reveal_type(test.get(good_keys))                  # N: Revealed type is "Union[__main__.A, __main__.B, None]"
-reveal_type(test.get(good_keys, 3))               # N: Revealed type is "Union[__main__.A, Literal[3]?, __main__.B]"
+reveal_type(test.get(good_keys))                  # N: Revealed type is "Union[__main__.A, __main__.B]"
+reveal_type(test.get(good_keys, 3))               # N: Revealed type is "Union[__main__.A, __main__.B]"
 reveal_type(test.pop(optional_keys))              # N: Revealed type is "Union[__main__.D, __main__.E]"
 reveal_type(test.pop(optional_keys, 3))           # N: Revealed type is "Union[__main__.D, __main__.E, Literal[3]?]"
 reveal_type(test.setdefault(good_keys, AAndB()))  # N: Revealed type is "Union[__main__.A, __main__.B]"
@@ -2037,15 +2037,18 @@ class D2(TypedDict):
     d: D
 
 x: Union[D1, D2]
-bad_keys: Literal['a', 'b', 'c', 'd']
 good_keys: Literal['b', 'c']
+mixed_keys: Literal['a', 'b', 'c', 'd']
+bad_keys: Literal['e', 'f']
 
-x[bad_keys]         # E: TypedDict "D1" has no key "d" \
+x[mixed_keys]         # E: TypedDict "D1" has no key "d" \
                     # E: TypedDict "D2" has no key "a"
 
 reveal_type(x[good_keys])           # N: Revealed type is "Union[__main__.B, __main__.C]"
-reveal_type(x.get(good_keys))       # N: Revealed type is "Union[__main__.B, __main__.C, None]"
-reveal_type(x.get(good_keys, 3))    # N: Revealed type is "Union[__main__.B, Literal[3]?, __main__.C]"
+reveal_type(x.get(good_keys))       # N: Revealed type is "Union[__main__.B, __main__.C]"
+reveal_type(x.get(good_keys, 3))    # N: Revealed type is "Union[__main__.B, __main__.C]"
+reveal_type(x.get(mixed_keys))      # N: Revealed type is "builtins.object"
+reveal_type(x.get(mixed_keys, 3))   # N: Revealed type is "builtins.object"
 reveal_type(x.get(bad_keys))        # N: Revealed type is "builtins.object"
 reveal_type(x.get(bad_keys, 3))     # N: Revealed type is "builtins.object"
 
diff --git a/test-data/unit/check-recursive-types.test b/test-data/unit/check-recursive-types.test
@@ -690,10 +690,11 @@ class TD(TypedDict, total=False):
     y: TD
 
 td: TD
+reveal_type(td.get("y"))  # N: Revealed type is "Union[TypedDict('__main__.TD', {'x'?: builtins.int, 'y'?: ...}), None]"
 td["y"] = {"x": 0, "y": {}}
 td["y"] = {"x": 0, "y": {"x": 0, "y": 42}}  # E: Incompatible types (expression has type "int", TypedDict item "y" has type "TD")
 
-reveal_type(td.get("y"))  # N: Revealed type is "Union[TypedDict('__main__.TD', {'x'?: builtins.int, 'y'?: TypedDict('__main__.TD', {'x'?: builtins.int, 'y'?: ...})}), None]"
+reveal_type(td.get("y"))  # N: Revealed type is "Union[TypedDict('__main__.TD', {'x'?: builtins.int, 'y'?: ...}), None]"
 s: str = td.get("y")  # E: Incompatible types in assignment (expression has type "Optional[TD]", variable has type "str")
 
 td.update({"x": 0, "y": {"x": 1, "y": {}}})
diff --git a/test-data/unit/check-typeddict.test b/test-data/unit/check-typeddict.test
diff --git a/test-data/unit/pythoneval.test b/test-data/unit/pythoneval.test
