Skip to content

More precise return types for TypedDict.get #19897

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 15 commits into from
Oct 14, 2025
Merged

More precise return types for TypedDict.get #19897

Show file tree
Hide file tree
Changes from 12 commits
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
27e50e2
improved TypedDict.get inference
randolf-scholz Aug 11, 2025
e708813
use original DictExpr test isntead of type-based
randolf-scholz Sep 21, 2025
63224f2
removed second overload
randolf-scholz Sep 22, 2025
808f593
added unit test for 19902
randolf-scholz Sep 22, 2025
de8effe
Clarify overload comments in checkmember.py
randolf-scholz Sep 23, 2025
4ebe192
check total/non-toal against union of keys
randolf-scholz Sep 23, 2025
a9d2388
added incremental test
randolf-scholz Sep 23, 2025
16c97ba
use appropriate namespace for TypedDict.get variable
randolf-scholz Sep 23, 2025
628a0be
Merge branch 'master' into special_case_typeddict_get
randolf-scholz Sep 23, 2025
443f6c7
Merge branch 'master' into special_case_typeddict_get
randolf-scholz Sep 27, 2025
6bf9401
Merge branch 'master' into special_case_typeddict_get
randolf-scholz Oct 4, 2025
cbbbd4d
Merge branch 'master' into special_case_typeddict_get
randolf-scholz Oct 11, 2025
63f4975
revert reveal_type changes
randolf-scholz Oct 12, 2025
f3e8556
Merge branch 'master' into special_case_typeddict_get
randolf-scholz Oct 12, 2025
a0fbcee
Remove extra whitespace
ilevkivskyi Oct 14, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
72 changes: 72 additions & 0 deletions mypy/checkmember.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -68,6 +69,7 @@
TypedDictType,
TypeOfAny,
TypeType,
TypeVarId,
TypeVarLikeType,
TypeVarTupleType,
TypeVarType,
Expand Down Expand Up @@ -1404,6 +1406,76 @@ def analyze_typeddict_access(
fallback=mx.chk.named_type("builtins.function"),
name=name,
)
elif name == "get":
Copy link
Member

@ilevkivskyi ilevkivskyi Oct 11, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am actually worried about the performance implications of this part of the PR. Some people have TypedDicts with few hundreds keys (not an exaggeration). This will then create an enormous overload. And overloads are slow to type-check (for various reasons). What exactly do we need this part for? Is it just for the better reveal_type(), or also for something else?

Copy link
Contributor Author

@randolf-scholz randolf-scholz Oct 11, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, I would say that is mostly for reveal_type. It will also change the error message if a user passes illegal arguments, but I don't think that is as important.

Copy link
Contributor Author

@randolf-scholz randolf-scholz Oct 11, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe it would be possible to:

  1. cache the overloads on the typeddict type itself (or make it a lazy property)
  2. Call the plugin early to skip all the overload checks.

Alternative, I could roll back this part of the PR and only keep the updated logic in the default plugin.

Copy link
Member

@ilevkivskyi ilevkivskyi Oct 12, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Calling existing plugin hook early would be a breaking change in the plugin system (IIRC we pass "default return type" as part of the plugin context), and adding a new hook is too big change for something like this. FWIW I don't think reveal_type() is really worth the potential problems, so I think it would be best to roll back this part.

Copy link
Contributor Author

@randolf-scholz randolf-scholz Oct 12, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There does seem to be one test where it makes a difference: (#19902)

# https://github.com/python/mypy/issues/19902
from typing import TypedDict, Union
from typing_extensions import TypeAlias, NotRequired
class A(TypedDict):
    key: NotRequired[int]

class B(TypedDict):
    key: NotRequired[int]

class C(TypedDict):
    key: NotRequired[int]

A_or_B: TypeAlias = Union[A, B]
A_or_B_or_C: TypeAlias = Union[A_or_B, C]

def test(d: A_or_B_or_C) -> None:
    reveal_type(d.get("key"))  # N: Revealed type is "Union[builtins.int, None]"

(on master this just gives object)

I think this should be fixable by doing a flatten_nested_union inside check_union_call_expr.

# 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_info = typ.fallback.type
# type variable for default value
tvar = TypeVarType(
Copy link
Collaborator

@sterliakov sterliakov Sep 23, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not sure that TypeVarType should ever be introduced without a corresponding symbol (TypeVarExpr). Any other place that needs to synthesize persistent (not single-use like a fake constructor in checkexpr) typevars does so in two steps (grep for TypeVarExpr in semanal_namedtuple.py or dataclasses plugin). I suspect that some nasal demons are expected if TypeVarType isn't anchored to some real symbol, this might bite us later.

"T",
f"{type_info.fullname}.get.T",
id=TypeVarId(-1, namespace=f"{type_info.fullname}.get"),
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.
# def (K, object=...) -> V
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 (implicit default)
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)

# def [T](K, T) -> V | T (explicit default)
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
# TODO: add support for extra items (PEP 728)
# 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)


Expand Down
80 changes: 25 additions & 55 deletions mypy/plugins/default.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 DictExpr, Expression, IntExpr, StrExpr, UnaryExpr
from mypy.plugin import (
AttributeContext,
ClassDefContext,
Expand Down Expand Up @@ -120,9 +120,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
Expand Down Expand Up @@ -212,46 +212,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 (
Expand All @@ -263,30 +223,40 @@ def typed_dict_get_callback(ctx: MethodContext) -> Type:
if keys is None:
return ctx.default_return_type

default_type: Type
default_arg: Expression | None
if len(ctx.arg_types) <= 1 or not ctx.arg_types[1]:
default_arg = None
default_type = NoneType()
elif len(ctx.arg_types[1]) == 1 and len(ctx.args[1]) == 1:
default_arg = ctx.args[1][0]
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, {})
if (
isinstance(default_arg, DictExpr)
and len(default_arg.items) == 0
and isinstance(value_type, TypedDictType)
and isinstance(vt := get_proper_type(value_type), TypedDictType)
):
# 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

Expand Down
Loading