From 27e50e25e34971736cb495359c47755b917afafd Mon Sep 17 00:00:00 2001 From: Randolf Scholz Date: Mon, 11 Aug 2025 14:30:04 +0200 Subject: [PATCH 01/10] improved TypedDict.get inference --- mypy/checkmember.py | 83 +++++++++++ mypy/plugins/default.py | 86 ++++------- test-data/unit/check-literal.test | 19 +-- test-data/unit/check-recursive-types.test | 3 +- test-data/unit/check-typeddict.test | 166 +++++++++++++++++++--- test-data/unit/pythoneval.test | 59 +++++--- 6 files changed, 315 insertions(+), 101 deletions(-) diff --git a/mypy/checkmember.py b/mypy/checkmember.py index f19a76ec6a34..c2ea77138b6f 100644 --- 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 index e492b8dd7335..ea73da9888fb 100644 --- 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 index 3c9290b8dbbb..ce0ae2844bae 100644 --- 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 index 86e9f02b5263..4c4d281418f9 100644 --- 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 index 34cae74d795b..2049ab738df8 100644 --- a/test-data/unit/check-typeddict.test +++ b/test-data/unit/check-typeddict.test @@ -997,10 +997,32 @@ if int(): -- Other TypedDict methods + +[case testTypedDictGetMethodOverloads] +from typing import TypedDict +from typing_extensions import Required, NotRequired + +class D(TypedDict): + a: int + b: NotRequired[str] + +def test(d: D) -> None: + reveal_type(d.get) # N: Revealed type is \ + "Overload(\ + def (Literal['a'], builtins.object =) -> builtins.int, \ + def (Literal['b']) -> Union[builtins.str, None], \ + def (Literal['b'], builtins.str) -> builtins.str, \ + def [T] (Literal['b'], T`-1) -> Union[builtins.str, T`-1], \ + def (builtins.str, builtins.object =) -> builtins.object)" + +[builtins fixtures/dict.pyi] +[typing fixtures/typing-typeddict.pyi] + + [case testTypedDictGetMethod] from typing import TypedDict class A: pass -D = TypedDict('D', {'x': int, 'y': str}) +D = TypedDict('D', {'x': int, 'y': str}, total=False) d: D reveal_type(d.get('x')) # N: Revealed type is "Union[builtins.int, None]" reveal_type(d.get('y')) # N: Revealed type is "Union[builtins.str, None]" @@ -1010,13 +1032,27 @@ reveal_type(d.get('y', None)) # N: Revealed type is "Union[builtins.str, None]" [builtins fixtures/dict.pyi] [typing fixtures/typing-typeddict.pyi] +[case testTypedDictGetMethod2] +from typing import TypedDict +class A: pass +D = TypedDict('D', {'x': int, 'y': str}, total=True) +d: D +reveal_type(d.get('x')) # N: Revealed type is "builtins.int" +reveal_type(d.get('y')) # N: Revealed type is "builtins.str" +reveal_type(d.get('x', A())) # N: Revealed type is "builtins.int" +reveal_type(d.get('x', 1)) # N: Revealed type is "builtins.int" +reveal_type(d.get('y', None)) # N: Revealed type is "builtins.str" +[builtins fixtures/dict.pyi] +[typing fixtures/typing-typeddict.pyi] + + [case testTypedDictGetMethodTypeContext] from typing import List, TypedDict class A: pass -D = TypedDict('D', {'x': List[int], 'y': int}) +D = TypedDict('D', {'x': List[int], 'y': int}, total=False) 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 +1062,16 @@ 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'], object = ..., /) -> int \ + # N: def get(Literal['y'], object = ..., /) -> str \ + # 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'], object = ..., /) -> int \ + # N: def get(Literal['y'], object = ..., /) -> str \ + # N: def get(str, object = ..., /) -> object x = d.get('z') reveal_type(x) # N: Revealed type is "builtins.object" s = '' @@ -1069,19 +1107,113 @@ p.get('x', 1 + 'y') # E: Unsupported operand types for + ("int" and "str") [case testTypedDictChainedGetWithEmptyDictDefault] from typing import TypedDict -C = TypedDict('C', {'a': int}) -D = TypedDict('D', {'x': C, 'y': str}) +C = TypedDict('C', {'a': int}, total=True) +D = TypedDict('D', {'x': C, 'y': str}, total=False) d: D -reveal_type(d.get('x', {})) \ - # N: Revealed type is "TypedDict('__main__.C', {'a'?: builtins.int})" -reveal_type(d.get('x', None)) \ - # N: Revealed type is "Union[TypedDict('__main__.C', {'a': builtins.int}), None]" +reveal_type(d.get('x', {})) # N: Revealed type is "TypedDict('__main__.C', {'a'?: builtins.int})" +reveal_type(d.get('x', None)) # N: Revealed type is "Union[TypedDict('__main__.C', {'a': builtins.int}), None]" reveal_type(d.get('x', {}).get('a')) # N: Revealed type is "Union[builtins.int, None]" reveal_type(d.get('x', {})['a']) # N: Revealed type is "builtins.int" [builtins fixtures/dict.pyi] [typing fixtures/typing-typeddict.pyi] +[case testTypedDictChainedGetWithEmptyDictDefault2] +from typing import TypedDict +C = TypedDict('C', {'a': int}, total=False) +D = TypedDict('D', {'x': C, 'y': str}, total=True) +d: D +reveal_type(d.get('x', {})) # N: Revealed type is "TypedDict('__main__.C', {'a'?: builtins.int})" +reveal_type(d.get('x', None)) # N: Revealed type is "TypedDict('__main__.C', {'a'?: builtins.int})" +reveal_type(d.get('x', {}).get('a')) # N: Revealed type is "Union[builtins.int, None]" +reveal_type(d.get('x', {})['a']) # N: Revealed type is "builtins.int" +[builtins fixtures/dict.pyi] +[typing fixtures/typing-typeddict.pyi] + + +[case testTypedDictChainedGetWithEmptyDictDefault3] +from typing import TypedDict +C = TypedDict('C', {'a': int}, total=True) +D = TypedDict('D', {'x': C, 'y': str}, total=True) +d: D +reveal_type(d.get('x', {})) # N: Revealed type is "TypedDict('__main__.C', {'a': builtins.int})" +reveal_type(d.get('x', None)) # N: Revealed type is "TypedDict('__main__.C', {'a': builtins.int})" +reveal_type(d.get('x', {}).get('a')) # N: Revealed type is "builtins.int" +reveal_type(d.get('x', {})['a']) # N: Revealed type is "builtins.int" +[builtins fixtures/dict.pyi] +[typing fixtures/typing-typeddict.pyi] + + +[case testTypedDictChainedGetWithEmptyDictDefault4] +from typing import TypedDict +C = TypedDict('C', {'a': int}, total=False) +D = TypedDict('D', {'x': C, 'y': str}, total=False) +d: D +reveal_type(d.get('x', {})) # N: Revealed type is "TypedDict('__main__.C', {'a'?: builtins.int})" +reveal_type(d.get('x', None)) # N: Revealed type is "Union[TypedDict('__main__.C', {'a'?: builtins.int}), None]" +reveal_type(d.get('x', {}).get('a')) # N: Revealed type is "Union[builtins.int, None]" +reveal_type(d.get('x', {})['a']) # N: Revealed type is "builtins.int" +[builtins fixtures/dict.pyi] +[typing fixtures/typing-typeddict.pyi] + + +[case testTypedDictGetMethodChained] +# check that chaining with get like ``.get(key, {}).get(subkey, {})`` works. +from typing import TypedDict, Mapping +from typing_extensions import Required, NotRequired, Never + +class Total(TypedDict, total=True): # no keys optional + key_one: int + key_two: str + +class Maybe(TypedDict, total=False): # all keys are optional + key_one: int + key_two: str + +class Mixed(TypedDict): # some keys optional + key_one: Required[int] + key_two: NotRequired[str] + +class Config(TypedDict): + required_total: Required[Total] + optional_total: NotRequired[Total] + required_mixed: Required[Mixed] + optional_mixed: NotRequired[Mixed] + required_maybe: Required[Maybe] + optional_maybe: NotRequired[Maybe] + +def test_chaining(d: Config) -> None: + reveal_type( d.get("required_total", {}) ) # N: Revealed type is "TypedDict('__main__.Total', {'key_one': builtins.int, 'key_two': builtins.str})" + reveal_type( d.get("optional_total", {}) ) # N: Revealed type is "TypedDict('__main__.Total', {'key_one'?: builtins.int, 'key_two'?: builtins.str})" + reveal_type( d.get("required_maybe", {}) ) # N: Revealed type is "TypedDict('__main__.Maybe', {'key_one'?: builtins.int, 'key_two'?: builtins.str})" + reveal_type( d.get("optional_maybe", {}) ) # N: Revealed type is "TypedDict('__main__.Maybe', {'key_one'?: builtins.int, 'key_two'?: builtins.str})" + reveal_type( d.get("required_mixed", {}) ) # N: Revealed type is "TypedDict('__main__.Mixed', {'key_one': builtins.int, 'key_two'?: builtins.str})" + reveal_type( d.get("optional_mixed", {}) ) # N: Revealed type is "TypedDict('__main__.Mixed', {'key_one'?: builtins.int, 'key_two'?: builtins.str})" + + reveal_type( d.get("required_total", {}).get("key_one") ) # N: Revealed type is "builtins.int" + reveal_type( d.get("required_total", {}).get("key_two") ) # N: Revealed type is "builtins.str" + reveal_type( d.get("required_total", {}).get("bad_key") ) # N: Revealed type is "builtins.object" + reveal_type( d.get("optional_total", {}).get("key_one") ) # N: Revealed type is "Union[builtins.int, None]" + reveal_type( d.get("optional_total", {}).get("key_two") ) # N: Revealed type is "Union[builtins.str, None]" + reveal_type( d.get("optional_total", {}).get("bad_key") ) # N: Revealed type is "builtins.object" + + reveal_type( d.get("required_maybe", {}).get("key_one") ) # N: Revealed type is "Union[builtins.int, None]" + reveal_type( d.get("required_maybe", {}).get("key_two") ) # N: Revealed type is "Union[builtins.str, None]" + reveal_type( d.get("required_maybe", {}).get("bad_key") ) # N: Revealed type is "builtins.object" + reveal_type( d.get("optional_maybe", {}).get("key_one") ) # N: Revealed type is "Union[builtins.int, None]" + reveal_type( d.get("optional_maybe", {}).get("key_two") ) # N: Revealed type is "Union[builtins.str, None]" + reveal_type( d.get("optional_maybe", {}).get("bad_key") ) # N: Revealed type is "builtins.object" + + reveal_type( d.get("required_mixed", {}).get("key_one") ) # N: Revealed type is "builtins.int" + reveal_type( d.get("required_mixed", {}).get("key_two") ) # N: Revealed type is "Union[builtins.str, None]" + reveal_type( d.get("required_mixed", {}).get("bad_key") ) # N: Revealed type is "builtins.object" + reveal_type( d.get("optional_mixed", {}).get("key_one") ) # N: Revealed type is "Union[builtins.int, None]" + reveal_type( d.get("optional_mixed", {}).get("key_two") ) # N: Revealed type is "Union[builtins.str, None]" + reveal_type( d.get("optional_mixed", {}).get("bad_key") ) # N: Revealed type is "builtins.object" +[builtins fixtures/dict.pyi] +[typing fixtures/typing-typeddict.pyi] + + -- Totality (the "total" keyword argument) [case testTypedDictWithTotalTrue] @@ -1769,8 +1901,8 @@ class TDB(TypedDict): td: Union[TDA, TDB] -reveal_type(td.get('a')) # N: Revealed type is "Union[builtins.int, None]" -reveal_type(td.get('b')) # N: Revealed type is "Union[builtins.str, None, builtins.int]" +reveal_type(td.get('a')) # N: Revealed type is "builtins.int" +reveal_type(td.get('b')) # N: Revealed type is "Union[builtins.str, builtins.int]" reveal_type(td.get('c')) # N: Revealed type is "builtins.object" reveal_type(td['a']) # N: Revealed type is "builtins.int" diff --git a/test-data/unit/pythoneval.test b/test-data/unit/pythoneval.test index 93b67bfa813a..efbe7645cecf 100644 --- a/test-data/unit/pythoneval.test +++ b/test-data/unit/pythoneval.test @@ -1034,24 +1034,47 @@ _program.py:17: note: Revealed type is "builtins.str" # Test that TypedDict get plugin works with typeshed stubs from typing import TypedDict class A: pass -D = TypedDict('D', {'x': int, 'y': str}) -d: D -reveal_type(d.get('x')) -reveal_type(d.get('y')) -reveal_type(d.get('z')) -d.get() -s = '' -reveal_type(d.get(s)) -[out] -_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: 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:11: note: Revealed type is "builtins.object" +D_total = TypedDict('D_total', {'x': int, 'y': str}, total=True) +D_not_total = TypedDict('D_not_total', {'x': int, 'y': str}, total=False) + +def test_total(d: D_total) -> None: + reveal_type(d.get('x')) + reveal_type(d.get('y')) + reveal_type(d.get('z')) + d.get() + s = '' + reveal_type(d.get(s)) + +def test_not_total(d: D_not_total) -> None: + reveal_type(d.get('x')) + reveal_type(d.get('y')) + reveal_type(d.get('z')) + d.get() + s = '' + reveal_type(d.get(s)) +[out] +_testTypedDictGet.py:8: note: Revealed type is "builtins.int" +_testTypedDictGet.py:9: note: Revealed type is "builtins.str" +_testTypedDictGet.py:10: note: Revealed type is "builtins.object" +_testTypedDictGet.py:11: error: All overload variants of "get" require at least one argument +_testTypedDictGet.py:11: note: Possible overload variants: +_testTypedDictGet.py:11: note: def get(Literal['x'], object = ..., /) -> int +_testTypedDictGet.py:11: note: def get(Literal['y'], object = ..., /) -> str +_testTypedDictGet.py:11: note: def get(str, object = ..., /) -> object +_testTypedDictGet.py:13: note: Revealed type is "builtins.object" +_testTypedDictGet.py:16: note: Revealed type is "Union[builtins.int, None]" +_testTypedDictGet.py:17: note: Revealed type is "Union[builtins.str, None]" +_testTypedDictGet.py:18: note: Revealed type is "builtins.object" +_testTypedDictGet.py:19: error: All overload variants of "get" require at least one argument +_testTypedDictGet.py:19: note: Possible overload variants: +_testTypedDictGet.py:19: note: def get(Literal['x'], /) -> Optional[int] +_testTypedDictGet.py:19: note: def get(Literal['x'], int, /) -> int +_testTypedDictGet.py:19: note: def [T] get(Literal['x'], T, /) -> Union[int, T] +_testTypedDictGet.py:19: note: def get(Literal['y'], /) -> Optional[str] +_testTypedDictGet.py:19: note: def get(Literal['y'], str, /) -> str +_testTypedDictGet.py:19: note: def [T] get(Literal['y'], T, /) -> Union[str, T] +_testTypedDictGet.py:19: note: def get(str, object = ..., /) -> object +_testTypedDictGet.py:21: note: Revealed type is "builtins.object" [case testTypedDictMappingMethods] from typing import TypedDict From e708813b3ec36290646da2439b9c58e4365f92ae Mon Sep 17 00:00:00 2001 From: Randolf Scholz Date: Sun, 21 Sep 2025 15:14:38 +0200 Subject: [PATCH 02/10] use original DictExpr test isntead of type-based --- mypy/plugins/default.py | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/mypy/plugins/default.py b/mypy/plugins/default.py index ea73da9888fb..ce4bac6d371f 100644 --- 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 IntExpr, StrExpr, UnaryExpr +from mypy.nodes import DictExpr, Expression, IntExpr, StrExpr, UnaryExpr from mypy.plugin import ( AttributeContext, ClassDefContext, @@ -75,7 +75,6 @@ TypedDictType, TypeOfAny, TypeVarType, - UninhabitedType, UnionType, get_proper_type, get_proper_types, @@ -225,9 +224,12 @@ def typed_dict_get_callback(ctx: MethodContext) -> Type: 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 @@ -242,14 +244,10 @@ def typed_dict_get_callback(ctx: MethodContext) -> Type: output_types.append(value_type) else: # 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) + isinstance(default_arg, DictExpr) + and len(default_arg.items) == 0 + and isinstance(vt := get_proper_type(value_type), TypedDictType) ): output_types.append(vt.copy_modified(required_keys=set())) else: From 63224f28c6e51f2793a52b02c5593ca760965560 Mon Sep 17 00:00:00 2001 From: Randolf Scholz Date: Mon, 22 Sep 2025 09:48:01 +0200 Subject: [PATCH 03/10] removed second overload --- mypy/checkmember.py | 13 ------------- test-data/unit/check-typeddict.test | 1 - test-data/unit/pythoneval.test | 2 -- 3 files changed, 16 deletions(-) diff --git a/mypy/checkmember.py b/mypy/checkmember.py index c2ea77138b6f..ce9a6b4d4a8d 100644 --- a/mypy/checkmember.py +++ b/mypy/checkmember.py @@ -1448,19 +1448,6 @@ def analyze_typeddict_access( ) 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], diff --git a/test-data/unit/check-typeddict.test b/test-data/unit/check-typeddict.test index 2049ab738df8..62ad0a7d894d 100644 --- a/test-data/unit/check-typeddict.test +++ b/test-data/unit/check-typeddict.test @@ -1011,7 +1011,6 @@ def test(d: D) -> None: "Overload(\ def (Literal['a'], builtins.object =) -> builtins.int, \ def (Literal['b']) -> Union[builtins.str, None], \ - def (Literal['b'], builtins.str) -> builtins.str, \ def [T] (Literal['b'], T`-1) -> Union[builtins.str, T`-1], \ def (builtins.str, builtins.object =) -> builtins.object)" diff --git a/test-data/unit/pythoneval.test b/test-data/unit/pythoneval.test index efbe7645cecf..36808d95765b 100644 --- a/test-data/unit/pythoneval.test +++ b/test-data/unit/pythoneval.test @@ -1068,10 +1068,8 @@ _testTypedDictGet.py:18: note: Revealed type is "builtins.object" _testTypedDictGet.py:19: error: All overload variants of "get" require at least one argument _testTypedDictGet.py:19: note: Possible overload variants: _testTypedDictGet.py:19: note: def get(Literal['x'], /) -> Optional[int] -_testTypedDictGet.py:19: note: def get(Literal['x'], int, /) -> int _testTypedDictGet.py:19: note: def [T] get(Literal['x'], T, /) -> Union[int, T] _testTypedDictGet.py:19: note: def get(Literal['y'], /) -> Optional[str] -_testTypedDictGet.py:19: note: def get(Literal['y'], str, /) -> str _testTypedDictGet.py:19: note: def [T] get(Literal['y'], T, /) -> Union[str, T] _testTypedDictGet.py:19: note: def get(str, object = ..., /) -> object _testTypedDictGet.py:21: note: Revealed type is "builtins.object" From 808f593bcfd3e7ddde35f3b31be100da27bef400 Mon Sep 17 00:00:00 2001 From: Randolf Scholz Date: Mon, 22 Sep 2025 10:39:00 +0200 Subject: [PATCH 04/10] added unit test for 19902 --- test-data/unit/check-typeddict.test | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/test-data/unit/check-typeddict.test b/test-data/unit/check-typeddict.test index 62ad0a7d894d..30b01e3024ac 100644 --- a/test-data/unit/check-typeddict.test +++ b/test-data/unit/check-typeddict.test @@ -1213,6 +1213,27 @@ def test_chaining(d: Config) -> None: [typing fixtures/typing-typeddict.pyi] +[case testTypedDictGetWithNestedUnionOfTypedDicts] +# 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]" +[builtins fixtures/dict.pyi] +[typing fixtures/typing-typeddict.pyi] + -- Totality (the "total" keyword argument) [case testTypedDictWithTotalTrue] From de8effee5df5874758e87f52b970787c28407f1f Mon Sep 17 00:00:00 2001 From: Randolf Scholz Date: Tue, 23 Sep 2025 19:18:22 +0200 Subject: [PATCH 05/10] Clarify overload comments in checkmember.py Updated comments in checkmember.py for clarity on overload definitions. --- mypy/checkmember.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/mypy/checkmember.py b/mypy/checkmember.py index ce9a6b4d4a8d..d21ff4732b9b 100644 --- a/mypy/checkmember.py +++ b/mypy/checkmember.py @@ -1426,6 +1426,7 @@ def analyze_typeddict_access( 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], @@ -1437,7 +1438,7 @@ def analyze_typeddict_access( overloads.append(overload) else: # The key is not required, but if it is present, we know its type. - # def (K) -> V | None + # def (K) -> V | None (implicit default) overload = CallableType( arg_types=[key_type], arg_kinds=[ARG_POS], @@ -1448,7 +1449,7 @@ def analyze_typeddict_access( ) overloads.append(overload) - # fallback: def [T](K, T) -> V | T + # def [T](K, T) -> V | T (explicit default) overload = CallableType( variables=[tvar], arg_types=[key_type, tvar], @@ -1461,6 +1462,7 @@ def analyze_typeddict_access( 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], From 4ebe192b2c60817440f728082ea3c70864c17764 Mon Sep 17 00:00:00 2001 From: Randolf Scholz Date: Tue, 23 Sep 2025 20:49:12 +0200 Subject: [PATCH 06/10] check total/non-toal against union of keys --- test-data/unit/check-typeddict.test | 110 ++++++++++++++++++++++++++-- 1 file changed, 102 insertions(+), 8 deletions(-) diff --git a/test-data/unit/check-typeddict.test b/test-data/unit/check-typeddict.test index 30b01e3024ac..ed4e72baf12a 100644 --- a/test-data/unit/check-typeddict.test +++ b/test-data/unit/check-typeddict.test @@ -1018,29 +1018,123 @@ def test(d: D) -> None: [typing fixtures/typing-typeddict.pyi] -[case testTypedDictGetMethod] -from typing import TypedDict -class A: pass +[case testTypedDictGetMethodTotalFalse] +from typing import TypedDict, Literal +class Unrelated: pass D = TypedDict('D', {'x': int, 'y': str}, total=False) d: D +u: Unrelated +x: Literal['x'] +y: Literal['y'] +z: Literal['z'] +x_or_y: Literal['x', 'y'] +x_or_z: Literal['x', 'z'] +x_or_y_or_z: Literal['x', 'y', 'z'] + +# test with literal expression reveal_type(d.get('x')) # N: Revealed type is "Union[builtins.int, None]" reveal_type(d.get('y')) # N: Revealed type is "Union[builtins.str, None]" -reveal_type(d.get('x', A())) # N: Revealed type is "Union[builtins.int, __main__.A]" +reveal_type(d.get('z')) # N: Revealed type is "builtins.object" +reveal_type(d.get('x', u)) # N: Revealed type is "Union[builtins.int, __main__.Unrelated]" reveal_type(d.get('x', 1)) # N: Revealed type is "builtins.int" reveal_type(d.get('y', None)) # N: Revealed type is "Union[builtins.str, None]" + +# test with literal type / union of literal types with implicit default +reveal_type(d.get(x)) # N: Revealed type is "Union[builtins.int, None]" +reveal_type(d.get(y)) # N: Revealed type is "Union[builtins.str, None]" +reveal_type(d.get(z)) # N: Revealed type is "builtins.object" +reveal_type(d.get(x_or_y)) # N: Revealed type is "Union[builtins.int, builtins.str, None]" +reveal_type(d.get(x_or_z)) # N: Revealed type is "builtins.object" +reveal_type(d.get(x_or_y_or_z)) # N: Revealed type is "builtins.object" + +# test with literal type / union of literal types with explicit default +reveal_type(d.get(x, u)) # N: Revealed type is "Union[builtins.int, __main__.Unrelated]" +reveal_type(d.get(y, u)) # N: Revealed type is "Union[builtins.str, __main__.Unrelated]" +reveal_type(d.get(z, u)) # N: Revealed type is "builtins.object" +reveal_type(d.get(x_or_y, u)) # N: Revealed type is "Union[builtins.int, builtins.str, __main__.Unrelated]" +reveal_type(d.get(x_or_z, u)) # N: Revealed type is "builtins.object" +reveal_type(d.get(x_or_y_or_z, u)) # N: Revealed type is "builtins.object" [builtins fixtures/dict.pyi] [typing fixtures/typing-typeddict.pyi] -[case testTypedDictGetMethod2] -from typing import TypedDict -class A: pass +[case testTypedDictGetMethodTotalTrue] +from typing import TypedDict, Literal +class Unrelated: pass D = TypedDict('D', {'x': int, 'y': str}, total=True) d: D +u: Unrelated +x: Literal['x'] +y: Literal['y'] +z: Literal['z'] +x_or_y: Literal['x', 'y'] +x_or_z: Literal['x', 'z'] +x_or_y_or_z: Literal['x', 'y', 'z'] + +# test with literal expression reveal_type(d.get('x')) # N: Revealed type is "builtins.int" reveal_type(d.get('y')) # N: Revealed type is "builtins.str" -reveal_type(d.get('x', A())) # N: Revealed type is "builtins.int" +reveal_type(d.get('z')) # N: Revealed type is "builtins.object" +reveal_type(d.get('x', u)) # N: Revealed type is "builtins.int" reveal_type(d.get('x', 1)) # N: Revealed type is "builtins.int" reveal_type(d.get('y', None)) # N: Revealed type is "builtins.str" + +# test with literal type / union of literal types with implicit default +reveal_type(d.get(x)) # N: Revealed type is "builtins.int" +reveal_type(d.get(y)) # N: Revealed type is "builtins.str" +reveal_type(d.get(z)) # N: Revealed type is "builtins.object" +reveal_type(d.get(x_or_y)) # N: Revealed type is "Union[builtins.int, builtins.str]" +reveal_type(d.get(x_or_z)) # N: Revealed type is "builtins.object" +reveal_type(d.get(x_or_y_or_z)) # N: Revealed type is "builtins.object" + +# test with literal type / union of literal types with explicit default +reveal_type(d.get(x, u)) # N: Revealed type is "builtins.int" +reveal_type(d.get(y, u)) # N: Revealed type is "builtins.str" +reveal_type(d.get(z, u)) # N: Revealed type is "builtins.object" +reveal_type(d.get(x_or_y, u)) # N: Revealed type is "Union[builtins.int, builtins.str]" +reveal_type(d.get(x_or_z, u)) # N: Revealed type is "builtins.object" +reveal_type(d.get(x_or_y_or_z, u)) # N: Revealed type is "builtins.object" + +[builtins fixtures/dict.pyi] +[typing fixtures/typing-typeddict.pyi] + + +[case testTypedDictGetMethodTotalMixed] +from typing import TypedDict, Literal +from typing_extensions import Required, NotRequired +class Unrelated: pass +D = TypedDict('D', {'x': Required[int], 'y': NotRequired[str]}) +d: D +u: Unrelated +x: Literal['x'] +y: Literal['y'] +z: Literal['z'] +x_or_y: Literal['x', 'y'] +x_or_z: Literal['x', 'z'] +x_or_y_or_z: Literal['x', 'y', 'z'] + +# test with literal expression +reveal_type(d.get('x')) # N: Revealed type is "builtins.int" +reveal_type(d.get('y')) # N: Revealed type is "Union[builtins.str, None]" +reveal_type(d.get('z')) # N: Revealed type is "builtins.object" +reveal_type(d.get('x', u)) # N: Revealed type is "builtins.int" +reveal_type(d.get('x', 1)) # N: Revealed type is "builtins.int" +reveal_type(d.get('y', None)) # N: Revealed type is "Union[builtins.str, None]" + +# test with literal type / union of literal types with implicit default +reveal_type(d.get(x)) # N: Revealed type is "builtins.int" +reveal_type(d.get(y)) # N: Revealed type is "Union[builtins.str, None]" +reveal_type(d.get(z)) # N: Revealed type is "builtins.object" +reveal_type(d.get(x_or_y)) # N: Revealed type is "Union[builtins.int, builtins.str, None]" +reveal_type(d.get(x_or_z)) # N: Revealed type is "builtins.object" +reveal_type(d.get(x_or_y_or_z)) # N: Revealed type is "builtins.object" + +# test with literal type / union of literal types with explicit default +reveal_type(d.get(x, u)) # N: Revealed type is "builtins.int" +reveal_type(d.get(y, u)) # N: Revealed type is "Union[builtins.str, __main__.Unrelated]" +reveal_type(d.get(z, u)) # N: Revealed type is "builtins.object" +reveal_type(d.get(x_or_y, u)) # N: Revealed type is "Union[builtins.int, builtins.str, __main__.Unrelated]" +reveal_type(d.get(x_or_z, u)) # N: Revealed type is "builtins.object" +reveal_type(d.get(x_or_y_or_z, u)) # N: Revealed type is "builtins.object" [builtins fixtures/dict.pyi] [typing fixtures/typing-typeddict.pyi] From a9d23889b9e8657c0c6d7ca9046bc12436913f16 Mon Sep 17 00:00:00 2001 From: Randolf Scholz Date: Tue, 23 Sep 2025 21:18:27 +0200 Subject: [PATCH 07/10] added incremental test --- test-data/unit/check-incremental.test | 201 ++++++++++++++++++++++++++ 1 file changed, 201 insertions(+) diff --git a/test-data/unit/check-incremental.test b/test-data/unit/check-incremental.test index d9d78715b396..52ef5f06aa78 100644 --- a/test-data/unit/check-incremental.test +++ b/test-data/unit/check-incremental.test @@ -7290,3 +7290,204 @@ tmp/m.py:3: error: Argument 1 to "accept_int" has incompatible type "str"; expec [out2] tmp/n.py:3: note: Revealed type is "builtins.str" tmp/m.py:3: error: Argument 1 to "accept_int" has incompatible type "str"; expected "int" + + +[case testIncrementalTypedDictGetMethodTotalFalse] +import impl +[file lib.py] +from typing import TypedDict +class Unrelated: pass +D = TypedDict('D', {'x': int, 'y': str}, total=False) +[file impl.py] +pass +[file impl.py.2] +from typing import Literal +from lib import D, Unrelated +d: D +u: Unrelated +x: Literal['x'] +y: Literal['y'] +z: Literal['z'] +x_or_y: Literal['x', 'y'] +x_or_z: Literal['x', 'z'] +x_or_y_or_z: Literal['x', 'y', 'z'] + +# test with literal expression +reveal_type(d.get('x')) +reveal_type(d.get('y')) +reveal_type(d.get('z')) +reveal_type(d.get('x', u)) +reveal_type(d.get('x', 1)) +reveal_type(d.get('y', None)) + +# test with literal type / union of literal types with implicit default +reveal_type(d.get(x)) +reveal_type(d.get(y)) +reveal_type(d.get(z)) +reveal_type(d.get(x_or_y)) +reveal_type(d.get(x_or_z)) +reveal_type(d.get(x_or_y_or_z)) + +# test with literal type / union of literal types with explicit default +reveal_type(d.get(x, u)) +reveal_type(d.get(y, u)) +reveal_type(d.get(z, u)) +reveal_type(d.get(x_or_y, u)) +reveal_type(d.get(x_or_z, u)) +reveal_type(d.get(x_or_y_or_z, u)) +[builtins fixtures/dict.pyi] +[typing fixtures/typing-typeddict.pyi] +[out] +[out2] +tmp/impl.py:13: note: Revealed type is "Union[builtins.int, None]" +tmp/impl.py:14: note: Revealed type is "Union[builtins.str, None]" +tmp/impl.py:15: note: Revealed type is "builtins.object" +tmp/impl.py:16: note: Revealed type is "Union[builtins.int, lib.Unrelated]" +tmp/impl.py:17: note: Revealed type is "builtins.int" +tmp/impl.py:18: note: Revealed type is "Union[builtins.str, None]" +tmp/impl.py:21: note: Revealed type is "Union[builtins.int, None]" +tmp/impl.py:22: note: Revealed type is "Union[builtins.str, None]" +tmp/impl.py:23: note: Revealed type is "builtins.object" +tmp/impl.py:24: note: Revealed type is "Union[builtins.int, builtins.str, None]" +tmp/impl.py:25: note: Revealed type is "builtins.object" +tmp/impl.py:26: note: Revealed type is "builtins.object" +tmp/impl.py:29: note: Revealed type is "Union[builtins.int, lib.Unrelated]" +tmp/impl.py:30: note: Revealed type is "Union[builtins.str, lib.Unrelated]" +tmp/impl.py:31: note: Revealed type is "builtins.object" +tmp/impl.py:32: note: Revealed type is "Union[builtins.int, builtins.str, lib.Unrelated]" +tmp/impl.py:33: note: Revealed type is "builtins.object" +tmp/impl.py:34: note: Revealed type is "builtins.object" + +[case testIncrementalTypedDictGetMethodTotalTrue] +import impl +[file lib.py] +from typing import TypedDict +class Unrelated: pass +D = TypedDict('D', {'x': int, 'y': str}, total=True) +[file impl.py] +pass +[file impl.py.2] +from typing import Literal +from lib import D, Unrelated +d: D +u: Unrelated +x: Literal['x'] +y: Literal['y'] +z: Literal['z'] +x_or_y: Literal['x', 'y'] +x_or_z: Literal['x', 'z'] +x_or_y_or_z: Literal['x', 'y', 'z'] + +# test with literal expression +reveal_type(d.get('x')) +reveal_type(d.get('y')) +reveal_type(d.get('z')) +reveal_type(d.get('x', u)) +reveal_type(d.get('x', 1)) +reveal_type(d.get('y', None)) + +# test with literal type / union of literal types with implicit default +reveal_type(d.get(x)) +reveal_type(d.get(y)) +reveal_type(d.get(z)) +reveal_type(d.get(x_or_y)) +reveal_type(d.get(x_or_z)) +reveal_type(d.get(x_or_y_or_z)) + +# test with literal type / union of literal types with explicit default +reveal_type(d.get(x, u)) +reveal_type(d.get(y, u)) +reveal_type(d.get(z, u)) +reveal_type(d.get(x_or_y, u)) +reveal_type(d.get(x_or_z, u)) +reveal_type(d.get(x_or_y_or_z, u)) +[builtins fixtures/dict.pyi] +[typing fixtures/typing-typeddict.pyi] +[out] +[out2] +tmp/impl.py:13: note: Revealed type is "builtins.int" +tmp/impl.py:14: note: Revealed type is "builtins.str" +tmp/impl.py:15: note: Revealed type is "builtins.object" +tmp/impl.py:16: note: Revealed type is "builtins.int" +tmp/impl.py:17: note: Revealed type is "builtins.int" +tmp/impl.py:18: note: Revealed type is "builtins.str" +tmp/impl.py:21: note: Revealed type is "builtins.int" +tmp/impl.py:22: note: Revealed type is "builtins.str" +tmp/impl.py:23: note: Revealed type is "builtins.object" +tmp/impl.py:24: note: Revealed type is "Union[builtins.int, builtins.str]" +tmp/impl.py:25: note: Revealed type is "builtins.object" +tmp/impl.py:26: note: Revealed type is "builtins.object" +tmp/impl.py:29: note: Revealed type is "builtins.int" +tmp/impl.py:30: note: Revealed type is "builtins.str" +tmp/impl.py:31: note: Revealed type is "builtins.object" +tmp/impl.py:32: note: Revealed type is "Union[builtins.int, builtins.str]" +tmp/impl.py:33: note: Revealed type is "builtins.object" +tmp/impl.py:34: note: Revealed type is "builtins.object" + + +[case testIncrementalTypedDictGetMethodTotalMixed] +import impl +[file lib.py] +from typing import TypedDict +from typing_extensions import Required, NotRequired +class Unrelated: pass +D = TypedDict('D', {'x': Required[int], 'y': NotRequired[str]}) +[file impl.py] +pass +[file impl.py.2] +from typing import Literal +from lib import D, Unrelated +d: D +u: Unrelated +x: Literal['x'] +y: Literal['y'] +z: Literal['z'] +x_or_y: Literal['x', 'y'] +x_or_z: Literal['x', 'z'] +x_or_y_or_z: Literal['x', 'y', 'z'] + +# test with literal expression +reveal_type(d.get('x')) +reveal_type(d.get('y')) +reveal_type(d.get('z')) +reveal_type(d.get('x', u)) +reveal_type(d.get('x', 1)) +reveal_type(d.get('y', None)) + +# test with literal type / union of literal types with implicit default +reveal_type(d.get(x)) +reveal_type(d.get(y)) +reveal_type(d.get(z)) +reveal_type(d.get(x_or_y)) +reveal_type(d.get(x_or_z)) +reveal_type(d.get(x_or_y_or_z)) + +# test with literal type / union of literal types with explicit default +reveal_type(d.get(x, u)) +reveal_type(d.get(y, u)) +reveal_type(d.get(z, u)) +reveal_type(d.get(x_or_y, u)) +reveal_type(d.get(x_or_z, u)) +reveal_type(d.get(x_or_y_or_z, u)) +[builtins fixtures/dict.pyi] +[typing fixtures/typing-typeddict.pyi] +[out] +[out2] +tmp/impl.py:13: note: Revealed type is "builtins.int" +tmp/impl.py:14: note: Revealed type is "Union[builtins.str, None]" +tmp/impl.py:15: note: Revealed type is "builtins.object" +tmp/impl.py:16: note: Revealed type is "builtins.int" +tmp/impl.py:17: note: Revealed type is "builtins.int" +tmp/impl.py:18: note: Revealed type is "Union[builtins.str, None]" +tmp/impl.py:21: note: Revealed type is "builtins.int" +tmp/impl.py:22: note: Revealed type is "Union[builtins.str, None]" +tmp/impl.py:23: note: Revealed type is "builtins.object" +tmp/impl.py:24: note: Revealed type is "Union[builtins.int, builtins.str, None]" +tmp/impl.py:25: note: Revealed type is "builtins.object" +tmp/impl.py:26: note: Revealed type is "builtins.object" +tmp/impl.py:29: note: Revealed type is "builtins.int" +tmp/impl.py:30: note: Revealed type is "Union[builtins.str, lib.Unrelated]" +tmp/impl.py:31: note: Revealed type is "builtins.object" +tmp/impl.py:32: note: Revealed type is "Union[builtins.int, builtins.str, lib.Unrelated]" +tmp/impl.py:33: note: Revealed type is "builtins.object" +tmp/impl.py:34: note: Revealed type is "builtins.object" From 16c97baf847d96a1b583037c07a8ac931788a331 Mon Sep 17 00:00:00 2001 From: Randolf Scholz Date: Tue, 23 Sep 2025 21:21:26 +0200 Subject: [PATCH 08/10] use appropriate namespace for TypedDict.get variable --- mypy/checkmember.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/mypy/checkmember.py b/mypy/checkmember.py index d21ff4732b9b..c8063fbbc3ce 100644 --- a/mypy/checkmember.py +++ b/mypy/checkmember.py @@ -1409,12 +1409,12 @@ 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") - + type_info = typ.fallback.type # type variable for default value tvar = TypeVarType( "T", - "T", - id=TypeVarId(-1), + 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), From 63f49759f40316d2ef7ea1ad4c02199fa3bee92f Mon Sep 17 00:00:00 2001 From: Randolf Scholz Date: Sun, 12 Oct 2025 17:31:13 +0200 Subject: [PATCH 09/10] revert reveal_type changes --- mypy/checkexpr.py | 2 +- mypy/checkmember.py | 72 -------------------- mypy/plugins/default.py | 46 ++++++++++++- test-data/unit/check-typeddict.test | 36 ++++++---- test-data/unit/fixtures/typing-typeddict.pyi | 4 +- test-data/unit/pythoneval.test | 18 +++-- 6 files changed, 77 insertions(+), 101 deletions(-) diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py index b8f9bf087467..3eb54579a050 100644 --- a/mypy/checkexpr.py +++ b/mypy/checkexpr.py @@ -1502,7 +1502,7 @@ def check_call_expr_with_callee_type( def check_union_call_expr(self, e: CallExpr, object_type: UnionType, member: str) -> Type: """Type check calling a member expression where the base type is a union.""" res: list[Type] = [] - for typ in object_type.relevant_items(): + for typ in flatten_nested_unions(object_type.relevant_items()): # Member access errors are already reported when visiting the member expression. with self.msg.filter_errors(): item = analyze_member_access( diff --git a/mypy/checkmember.py b/mypy/checkmember.py index 7e33a34e8622..719b48b14e07 100644 --- a/mypy/checkmember.py +++ b/mypy/checkmember.py @@ -17,7 +17,6 @@ from mypy.meet import is_overlapping_types from mypy.messages import MessageBuilder from mypy.nodes import ( - ARG_OPT, ARG_POS, ARG_STAR, ARG_STAR2, @@ -69,7 +68,6 @@ TypedDictType, TypeOfAny, TypeType, - TypeVarId, TypeVarLikeType, TypeVarTupleType, TypeVarType, @@ -1406,76 +1404,6 @@ 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_info = typ.fallback.type - # type variable for default value - tvar = TypeVarType( - "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) diff --git a/mypy/plugins/default.py b/mypy/plugins/default.py index ce4bac6d371f..7a58307fc559 100644 --- a/mypy/plugins/default.py +++ b/mypy/plugins/default.py @@ -120,9 +120,9 @@ def get_function_signature_hook( def get_method_signature_hook( self, fullname: str ) -> Callable[[MethodSigContext], FunctionLike] | None: - # NOTE: signatures for `__setitem__`, `__delitem__` and `get` are - # defined in checkmember.py/analyze_typeddict_access - if fullname in TD_SETDEFAULT_NAMES: + if fullname == "typing.Mapping.get": + return typed_dict_get_signature_callback + elif fullname in TD_SETDEFAULT_NAMES: return typed_dict_setdefault_signature_callback elif fullname in TD_POP_NAMES: return typed_dict_pop_signature_callback @@ -212,6 +212,46 @@ 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 ( diff --git a/test-data/unit/check-typeddict.test b/test-data/unit/check-typeddict.test index ed4e72baf12a..e184dcc10169 100644 --- a/test-data/unit/check-typeddict.test +++ b/test-data/unit/check-typeddict.test @@ -1007,12 +1007,12 @@ class D(TypedDict): b: NotRequired[str] def test(d: D) -> None: - reveal_type(d.get) # N: Revealed type is \ - "Overload(\ - def (Literal['a'], builtins.object =) -> builtins.int, \ - def (Literal['b']) -> Union[builtins.str, None], \ - def [T] (Literal['b'], T`-1) -> Union[builtins.str, T`-1], \ - def (builtins.str, builtins.object =) -> builtins.object)" + reveal_type(d.get) # N: Revealed type is "Overload(def (k: builtins.str) -> builtins.object, def (builtins.str, builtins.object) -> builtins.object, def [V] (builtins.str, V`4) -> builtins.object)" + + + + + [builtins fixtures/dict.pyi] [typing fixtures/typing-typeddict.pyi] @@ -1155,21 +1155,29 @@ 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" require at least one argument \ +d.get() # E: All overload variants of "get" of "Mapping" require at least one argument \ # N: Possible overload variants: \ - # N: def get(Literal['x'], object = ..., /) -> int \ - # N: def get(Literal['y'], object = ..., /) -> str \ - # N: def get(str, object = ..., /) -> object -d.get('x', 1, 2) # E: No overload variant of "get" matches argument types "str", "int", "int" \ + # N: def get(self, k: str) -> object \ + # N: def get(self, str, object, /) -> object \ + # N: def [V] get(self, str, V, /) -> object +d.get('x', 1, 2) # E: No overload variant of "get" of "Mapping" matches argument types "str", "int", "int" \ # N: Possible overload variants: \ - # N: def get(Literal['x'], object = ..., /) -> int \ - # N: def get(Literal['y'], object = ..., /) -> str \ - # N: def get(str, object = ..., /) -> object + # N: def get(self, k: str) -> object \ + # N: def get(self, str, object, /) -> object \ + # N: def [V] get(self, str, Union[int, V], /) -> object x = d.get('z') reveal_type(x) # N: Revealed type is "builtins.object" s = '' y = d.get(s) reveal_type(y) # N: Revealed type is "builtins.object" + + + + + + + + [builtins fixtures/dict.pyi] [typing fixtures/typing-typeddict.pyi] diff --git a/test-data/unit/fixtures/typing-typeddict.pyi b/test-data/unit/fixtures/typing-typeddict.pyi index 16658c82528b..29635b651870 100644 --- a/test-data/unit/fixtures/typing-typeddict.pyi +++ b/test-data/unit/fixtures/typing-typeddict.pyi @@ -56,7 +56,9 @@ class Mapping(Iterable[T], Generic[T, T_co], metaclass=ABCMeta): @overload def get(self, k: T) -> Optional[T_co]: pass @overload - def get(self, k: T, default: Union[T_co, V]) -> Union[T_co, V]: pass + def get(self, k: T, default: T_co, /) -> Optional[T_co]: pass # type: ignore[misc] + @overload + def get(self, k: T, default: V, /) -> Union[T_co, V]: pass def values(self) -> Iterable[T_co]: pass # Approximate return type def __len__(self) -> int: ... def __contains__(self, arg: object) -> int: pass diff --git a/test-data/unit/pythoneval.test b/test-data/unit/pythoneval.test index 36808d95765b..2069d082df17 100644 --- a/test-data/unit/pythoneval.test +++ b/test-data/unit/pythoneval.test @@ -1056,22 +1056,20 @@ def test_not_total(d: D_not_total) -> None: _testTypedDictGet.py:8: note: Revealed type is "builtins.int" _testTypedDictGet.py:9: note: Revealed type is "builtins.str" _testTypedDictGet.py:10: note: Revealed type is "builtins.object" -_testTypedDictGet.py:11: error: All overload variants of "get" require at least one argument +_testTypedDictGet.py:11: error: All overload variants of "get" of "Mapping" require at least one argument _testTypedDictGet.py:11: note: Possible overload variants: -_testTypedDictGet.py:11: note: def get(Literal['x'], object = ..., /) -> int -_testTypedDictGet.py:11: note: def get(Literal['y'], object = ..., /) -> str -_testTypedDictGet.py:11: note: def get(str, object = ..., /) -> object +_testTypedDictGet.py:11: note: def get(self, str, /) -> object +_testTypedDictGet.py:11: note: def get(self, str, /, default: object) -> object +_testTypedDictGet.py:11: note: def [_T] get(self, str, /, default: _T) -> object _testTypedDictGet.py:13: note: Revealed type is "builtins.object" _testTypedDictGet.py:16: note: Revealed type is "Union[builtins.int, None]" _testTypedDictGet.py:17: note: Revealed type is "Union[builtins.str, None]" _testTypedDictGet.py:18: note: Revealed type is "builtins.object" -_testTypedDictGet.py:19: error: All overload variants of "get" require at least one argument +_testTypedDictGet.py:19: error: All overload variants of "get" of "Mapping" require at least one argument _testTypedDictGet.py:19: note: Possible overload variants: -_testTypedDictGet.py:19: note: def get(Literal['x'], /) -> Optional[int] -_testTypedDictGet.py:19: note: def [T] get(Literal['x'], T, /) -> Union[int, T] -_testTypedDictGet.py:19: note: def get(Literal['y'], /) -> Optional[str] -_testTypedDictGet.py:19: note: def [T] get(Literal['y'], T, /) -> Union[str, T] -_testTypedDictGet.py:19: note: def get(str, object = ..., /) -> object +_testTypedDictGet.py:19: note: def get(self, str, /) -> object +_testTypedDictGet.py:19: note: def get(self, str, /, default: object) -> object +_testTypedDictGet.py:19: note: def [_T] get(self, str, /, default: _T) -> object _testTypedDictGet.py:21: note: Revealed type is "builtins.object" [case testTypedDictMappingMethods] From a0fbcee4a2b5011cd15531ba62004a6947689d52 Mon Sep 17 00:00:00 2001 From: Ivan Levkivskyi Date: Tue, 14 Oct 2025 15:21:56 +0100 Subject: [PATCH 10/10] Remove extra whitespace --- test-data/unit/check-typeddict.test | 14 -------------- 1 file changed, 14 deletions(-) diff --git a/test-data/unit/check-typeddict.test b/test-data/unit/check-typeddict.test index e184dcc10169..e1a70efe9316 100644 --- a/test-data/unit/check-typeddict.test +++ b/test-data/unit/check-typeddict.test @@ -1008,12 +1008,6 @@ class D(TypedDict): def test(d: D) -> None: reveal_type(d.get) # N: Revealed type is "Overload(def (k: builtins.str) -> builtins.object, def (builtins.str, builtins.object) -> builtins.object, def [V] (builtins.str, V`4) -> builtins.object)" - - - - - - [builtins fixtures/dict.pyi] [typing fixtures/typing-typeddict.pyi] @@ -1170,14 +1164,6 @@ reveal_type(x) # N: Revealed type is "builtins.object" s = '' y = d.get(s) reveal_type(y) # N: Revealed type is "builtins.object" - - - - - - - - [builtins fixtures/dict.pyi] [typing fixtures/typing-typeddict.pyi]