improved TypedDict.get inference

Author: randolf-scholz

mypy/checkmember.py

@@ -1409,57 +1409,65 @@ def analyze_typeddict_access(
         str_type = mx.chk.named_type("builtins.str")
         fn_type = mx.chk.named_type("builtins.function")
         object_type = mx.chk.named_type("builtins.object")
-        t = TypeVarType(
+
+        # 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, val in typ.items.items():
+        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=[LiteralType(key, fallback=str_type), object_type],
+                    arg_types=[key_type, object_type],
                     arg_kinds=[ARG_POS, ARG_OPT],
                     arg_names=[None, None],
-                    ret_type=val,
+                    ret_type=value_type,
                     fallback=fn_type,
                     name=name,
                 )
                 overloads.append(overload)
             else:
-                # The key is not required, so we add the overloads:
-                # def (Literal[Key]) -> Val | None
-                # def (Literal[Key], default: Val) -> Val
-                # def [T] (Literal[Key], default: T = ..., /) -> Val | T
-                # TODO: simplify the last two overloads to just one
+                # The key is not required, but if it is present, we know its type.
+                # def (K) -> V | None
                 overload = CallableType(
-                    arg_types=[LiteralType(key, fallback=str_type)],
+                    arg_types=[key_type],
                     arg_kinds=[ARG_POS],
                     arg_names=[None],
-                    ret_type=UnionType.make_union([val, NoneType()]),
+                    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=[LiteralType(key, fallback=str_type), val],
+                    arg_types=[key_type, value_type],
                     arg_kinds=[ARG_POS, ARG_POS],
                     arg_names=[None, None],
-                    ret_type=val,
+                    ret_type=value_type,
                     fallback=fn_type,
                     name=name,
                 )
                 overloads.append(overload)
+
+                # fallback: def [T](K, T) -> V | T
                 overload = CallableType(
-                    variables=[t],
-                    arg_types=[LiteralType(key, fallback=str_type), t],
-                    arg_kinds=[ARG_POS, ARG_OPT],
+                    variables=[tvar],
+                    arg_types=[key_type, tvar],
+                    arg_kinds=[ARG_POS, ARG_POS],
                     arg_names=[None, None],
-                    ret_type=UnionType.make_union([val, t]),
+                    ret_type=UnionType.make_union([value_type, tvar]),
                     fallback=fn_type,
                     name=name,
                 )

mypy/plugins/default.py

@@ -70,12 +70,12 @@
     Instance,
     LiteralType,
     NoneType,
-    ProperType,
     TupleType,
     Type,
     TypedDictType,
     TypeOfAny,
     TypeVarType,
+    UninhabitedType,
     UnionType,
     get_proper_type,
     get_proper_types,
@@ -122,7 +122,7 @@ def get_method_signature_hook(
         self, fullname: str
     ) -> Callable[[MethodSigContext], FunctionLike] | None:
         # NOTE: signatures for `__setitem__`, `__delitem__` and `get` are
-        #  checkmember.py/analyze_typeddict_access
+        #  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:
@@ -214,41 +214,51 @@ def get_class_decorator_hook_2(
 
 
 def typed_dict_get_callback(ctx: MethodContext) -> Type:
-    """Infer a precise return type for TypedDict.get with union of literal first argument."""
+    """Infer a precise return type for TypedDict.get with literal first argument."""
     if (
         isinstance(ctx.type, TypedDictType)
         and len(ctx.arg_types) >= 1
         and len(ctx.arg_types[0]) == 1
-        and isinstance(get_proper_type(ctx.arg_types[0][0]), UnionType)
     ):
         keys = try_getting_str_literals(ctx.args[0][0], ctx.arg_types[0][0])
         if keys is None:
             return ctx.default_return_type
 
-        output_types: list[Type] = []
-        default_arg: Type
-
+        default_type: Type
         if len(ctx.arg_types) <= 1 or not ctx.arg_types[1]:
-            default_arg = NoneType()
+            default_type = NoneType()
         elif len(ctx.arg_types[1]) == 1 and len(ctx.args[1]) == 1:
-            default_arg = ctx.arg_types[1][0]
+            default_type = ctx.arg_types[1][0]
         else:
             return ctx.default_return_type
 
+        output_types: list[Type] = []
         for key in keys:
-            value: ProperType | None = get_proper_type(ctx.type.items.get(key))
-            if value is None:
+            value_type: Type | None = ctx.type.items.get(key)
+            if value_type is None:
                 return ctx.default_return_type
 
             if key in ctx.type.required_keys:
-                output_types.append(value)
+                output_types.append(value_type)
             else:
-                output_types.append(value)
-                output_types.append(default_arg)
+                # HACK to deal with get(key, {})
+                proper_default = get_proper_type(default_type)
+                if (
+                    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)
+                ):
+                    output_types.append(vt.copy_modified(required_keys=set()))
+                else:
+                    output_types.append(value_type)
+                    output_types.append(default_type)
 
         # for nicer reveal_type, put default at the end, if it is present
-        if default_arg in output_types:
-            output_types = [t for t in output_types if t != default_arg] + [default_arg]
+        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
 

test-data/unit/check-recursive-types.test

@@ -690,6 +690,7 @@ 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")