diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py index 6cdd64e84302f..3c8546816c3e0 100644 --- a/mypy/checkexpr.py +++ b/mypy/checkexpr.py @@ -694,6 +694,17 @@ def method_fullname(self, object_type: Type, method_name: str) -> str | None: elif isinstance(object_type, TypeType): object_type = object_type.item + # A TypeVar has no qualified name of its own. For `.get`, resolve a + # TypedDict upper bound through the same fallback owner as a direct + # TypedDict (typing.Mapping.get) so the existing hooks can run. The + # original receiver type is left unchanged outside this lookup. Other + # methods stay unnamed: their hooks assume a TypedDict receiver and + # would rewrite signatures for a TypeVar. Other bound kinds are unchanged. + if isinstance(object_type, TypeVarType) and method_name == "get": + bound = get_proper_type(object_type.upper_bound) + if isinstance(bound, TypedDictType): + object_type = bound + type_name = None if isinstance(object_type, Instance): type_name = object_type.type.fullname diff --git a/mypy/plugins/default.py b/mypy/plugins/default.py index 67352a05e4474..cb84074d01e2c 100644 --- a/mypy/plugins/default.py +++ b/mypy/plugins/default.py @@ -226,6 +226,20 @@ def len_callback(ctx: FunctionContext) -> Type: return ctx.default_return_type +def _typed_dict_get_receiver(typ: Type) -> TypedDictType | None: + """TypedDict for `.get`, including a TypeVar's proper upper bound. + + The plugin context still carries the original receiver. Callers use this + only for field lookup, required keys, and closedness. + """ + typ = get_proper_type(typ) + if isinstance(typ, TypeVarType): + typ = get_proper_type(typ.upper_bound) + if isinstance(typ, TypedDictType): + return typ + return None + + def typed_dict_get_signature_callback(ctx: MethodSigContext) -> CallableType: """Try to infer a better signature type for TypedDict.get. @@ -233,8 +247,9 @@ def typed_dict_get_signature_callback(ctx: MethodSigContext) -> CallableType: depends on a TypedDict value type. """ signature = ctx.default_signature + receiver = _typed_dict_get_receiver(ctx.type) if ( - isinstance(ctx.type, TypedDictType) + receiver is not None and len(ctx.args) == 2 and len(ctx.args[0]) == 1 and isinstance(ctx.args[0][0], StrExpr) @@ -243,7 +258,7 @@ def typed_dict_get_signature_callback(ctx: MethodSigContext) -> CallableType: and len(ctx.args[1]) == 1 ): key = ctx.args[0][0].value - value_type = get_proper_type(ctx.type.items.get(key)) + value_type = get_proper_type(receiver.items.get(key)) ret_type = signature.ret_type if value_type: default_arg = ctx.args[1][0] @@ -268,11 +283,8 @@ def typed_dict_get_signature_callback(ctx: MethodSigContext) -> CallableType: def typed_dict_get_callback(ctx: MethodContext) -> Type: """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 - ): + receiver = _typed_dict_get_receiver(ctx.type) + if receiver is not None and len(ctx.arg_types) >= 1 and len(ctx.arg_types[0]) == 1: keys = try_getting_str_literals(ctx.args[0][0], ctx.arg_types[0][0]) if keys is None: return ctx.default_return_type @@ -290,12 +302,12 @@ def typed_dict_get_callback(ctx: MethodContext) -> Type: output_types: list[Type] = [] for key in keys: - value_type: Type | None = ctx.type.items.get(key) + value_type: Type | None = receiver.items.get(key) if value_type is None: - if not ctx.type.is_closed: + if not receiver.is_closed: return ctx.default_return_type output_types.append(default_type) - elif key in ctx.type.required_keys: + elif key in receiver.required_keys: output_types.append(value_type) else: # HACK to deal with get(key, {}) diff --git a/test-data/unit/check-python312.test b/test-data/unit/check-python312.test index 9d32689a1e989..3fb3e230d92d7 100644 --- a/test-data/unit/check-python312.test +++ b/test-data/unit/check-python312.test @@ -2401,3 +2401,48 @@ reveal_type(baz(objs)) # N: Revealed type is "builtins.list[__main__.B | __main reveal_type(baz(objs_a)) # N: Revealed type is "builtins.list[__main__.B | __main__.C]" [builtins fixtures/tuple.pyi] [typing fixtures/typing-full.pyi] + +[case testPEP695TypedDictGetBound] +# https://github.com/python/mypy/issues/16872 +from typing import TypedDict + +class MyDict(TypedDict): + foo: list[str] + bar: int + +class Partial(TypedDict, total=False): + foo: list[str] + bar: int + +class Child(MyDict): + extra: int + +def required[T: MyDict](d: T, direct: MyDict) -> T: + reveal_type(d) # N: Revealed type is "T`-1" + reveal_type(d['foo']) # N: Revealed type is "builtins.list[builtins.str]" + reveal_type(direct['foo']) # N: Revealed type is "builtins.list[builtins.str]" + reveal_type(d.get('foo')) # N: Revealed type is "builtins.list[builtins.str]" + reveal_type(direct.get('foo')) # N: Revealed type is "builtins.list[builtins.str]" + reveal_type(d.get('foo', ['x'])) # N: Revealed type is "builtins.list[builtins.str]" + reveal_type(direct.get('foo', ['x'])) # N: Revealed type is "builtins.list[builtins.str]" + reveal_type(d.get('bar')) # N: Revealed type is "builtins.int" + reveal_type(direct.get('bar', 0)) # N: Revealed type is "builtins.int" + reveal_type(d.get('bar', 0)) # N: Revealed type is "builtins.int" + return d + +def optional[T: Partial](d: T, direct: Partial) -> T: + reveal_type(d.get('foo')) # N: Revealed type is "builtins.list[builtins.str] | None" + reveal_type(direct.get('foo')) # N: Revealed type is "builtins.list[builtins.str] | None" + reveal_type(d.get('bar', None)) # N: Revealed type is "builtins.int | None" + reveal_type(direct.get('bar', None)) # N: Revealed type is "builtins.int | None" + reveal_type(d.get('foo', [])) # N: Revealed type is "builtins.list[builtins.str]" + reveal_type(direct.get('foo', [])) # N: Revealed type is "builtins.list[builtins.str]" + reveal_type(d.get('bar', 1)) # N: Revealed type is "builtins.int" + reveal_type(direct.get('bar', 1)) # N: Revealed type is "builtins.int" + return d + +def use(child: Child, partial: Partial) -> None: + reveal_type(required(child, child)) # N: Revealed type is "TypedDict('__main__.Child', {'foo': builtins.list[builtins.str], 'bar': builtins.int, 'extra': builtins.int})" + reveal_type(optional(partial, partial)) # N: Revealed type is "TypedDict('__main__.Partial', {'foo'?: builtins.list[builtins.str], 'bar'?: builtins.int})" +[builtins fixtures/tuple.pyi] +[typing fixtures/typing-full.pyi] diff --git a/test-data/unit/check-typeddict.test b/test-data/unit/check-typeddict.test index 978f0d849c23b..460a674078dfd 100644 --- a/test-data/unit/check-typeddict.test +++ b/test-data/unit/check-typeddict.test @@ -6148,3 +6148,208 @@ def test(f: Callable[[object], object]) -> None: ... test(foo) # E: Argument 1 to "test" has incompatible type overloaded function; expected "Callable[[object], object]" [builtins fixtures/dict.pyi] [typing fixtures/typing-typeddict.pyi] + +[case testTypedDictGetBoundTypeVarRequired] +# https://github.com/python/mypy/issues/16872 +from typing import TypedDict, TypeVar + +class MyDict(TypedDict): + foo: list[str] + +class Child(MyDict): + extra: int + +class Unrelated: pass + +T = TypeVar('T', bound=MyDict) + +def f(d: T, direct: MyDict, u: Unrelated) -> T: + reveal_type(d) # N: Revealed type is "T`-1" + reveal_type(d['foo']) # N: Revealed type is "builtins.list[builtins.str]" + reveal_type(direct['foo']) # N: Revealed type is "builtins.list[builtins.str]" + reveal_type(d.get('foo')) # N: Revealed type is "builtins.list[builtins.str]" + reveal_type(direct.get('foo')) # N: Revealed type is "builtins.list[builtins.str]" + reveal_type(d.get('foo', ['x'])) # N: Revealed type is "builtins.list[builtins.str]" + reveal_type(direct.get('foo', ['x'])) # N: Revealed type is "builtins.list[builtins.str]" + reveal_type(d.get('foo', [])) # N: Revealed type is "builtins.list[builtins.str]" + reveal_type(direct.get('foo', [])) # N: Revealed type is "builtins.list[builtins.str]" + reveal_type(d.get('foo', u)) # N: Revealed type is "builtins.list[builtins.str]" + reveal_type(direct.get('foo', u)) # N: Revealed type is "builtins.list[builtins.str]" + return d + +def g(child: Child) -> None: + reveal_type(f(child, child, Unrelated())) # N: Revealed type is "TypedDict('__main__.Child', {'foo': builtins.list[builtins.str], 'extra': builtins.int})" +[builtins fixtures/dict.pyi] +[typing fixtures/typing-typeddict.pyi] + +[case testTypedDictGetBoundTypeVarOptional] +from typing import TypedDict, TypeVar +from typing_extensions import NotRequired, Required + +class Unrelated: pass + +class Partial(TypedDict, total=False): + x: int + y: str + +class Mixed(TypedDict): + x: Required[int] + y: NotRequired[str] + +P = TypeVar('P', bound=Partial) +M = TypeVar('M', bound=Mixed) + +def partial(d: P, direct: Partial, u: Unrelated) -> None: + reveal_type(d.get('x')) # N: Revealed type is "builtins.int | None" + reveal_type(direct.get('x')) # N: Revealed type is "builtins.int | None" + reveal_type(d.get('y', u)) # N: Revealed type is "builtins.str | __main__.Unrelated" + reveal_type(direct.get('y', u)) # N: Revealed type is "builtins.str | __main__.Unrelated" + reveal_type(d.get('x', 1)) # N: Revealed type is "builtins.int" + reveal_type(direct.get('x', 1)) # N: Revealed type is "builtins.int" + reveal_type(d.get('y', None)) # N: Revealed type is "builtins.str | None" + reveal_type(direct.get('y', None)) # N: Revealed type is "builtins.str | None" + +def mixed(d: M, direct: Mixed, u: Unrelated) -> None: + reveal_type(d.get('x')) # N: Revealed type is "builtins.int" + reveal_type(direct.get('x')) # N: Revealed type is "builtins.int" + reveal_type(d.get('y')) # N: Revealed type is "builtins.str | None" + reveal_type(direct.get('y')) # N: Revealed type is "builtins.str | None" + reveal_type(d.get('x', u)) # N: Revealed type is "builtins.int" + reveal_type(direct.get('x', u)) # N: Revealed type is "builtins.int" + reveal_type(d.get('y', u)) # N: Revealed type is "builtins.str | __main__.Unrelated" + reveal_type(direct.get('y', u)) # N: Revealed type is "builtins.str | __main__.Unrelated" + reveal_type(d.get('y', 'a')) # N: Revealed type is "builtins.str" + reveal_type(direct.get('y', 'a')) # N: Revealed type is "builtins.str" + reveal_type(d.get('x', 1)) # N: Revealed type is "builtins.int" + reveal_type(direct.get('x', 1)) # N: Revealed type is "builtins.int" +[builtins fixtures/dict.pyi] +[typing fixtures/typing-typeddict.pyi] + +[case testTypedDictGetBoundTypeVarTypeContext] +from typing import TypedDict, TypeVar + +class Inner(TypedDict): + a: int + +class Outer(TypedDict, total=False): + items: list[int] + nested: Inner + +class TotalOuter(TypedDict): + nested: Inner + items: list[int] + +T = TypeVar('T', bound=Outer) +U = TypeVar('U', bound=TotalOuter) + +def optional(d: T, direct: Outer) -> None: + reveal_type(d.get('items', [])) # N: Revealed type is "builtins.list[builtins.int]" + reveal_type(direct.get('items', [])) # N: Revealed type is "builtins.list[builtins.int]" + reveal_type(d.get('nested', {})) # N: Revealed type is "TypedDict('__main__.Inner', {'a'?: builtins.int})" + reveal_type(direct.get('nested', {})) # N: Revealed type is "TypedDict('__main__.Inner', {'a'?: builtins.int})" + reveal_type(d.get('nested', {}).get('a')) # N: Revealed type is "builtins.int | None" + reveal_type(direct.get('nested', {}).get('a')) # N: Revealed type is "builtins.int | None" + reveal_type(d.get('nested', {})['a']) # N: Revealed type is "builtins.int" + reveal_type(direct.get('nested', {})['a']) # N: Revealed type is "builtins.int" + +def required(d: U, direct: TotalOuter) -> None: + reveal_type(d.get('nested', {})) # N: Revealed type is "TypedDict('__main__.Inner', {'a': builtins.int})" + reveal_type(direct.get('nested', {})) # N: Revealed type is "TypedDict('__main__.Inner', {'a': builtins.int})" + reveal_type(d.get('nested', {}).get('a')) # N: Revealed type is "builtins.int" + reveal_type(direct.get('nested', {}).get('a')) # N: Revealed type is "builtins.int" + reveal_type(d.get('items', [])) # N: Revealed type is "builtins.list[builtins.int]" + reveal_type(direct.get('items', [])) # N: Revealed type is "builtins.list[builtins.int]" +[builtins fixtures/dict.pyi] +[typing fixtures/typing-typeddict.pyi] + +[case testTypedDictGetBoundTypeVarLiteralKeys] +from typing import Literal, TypedDict, TypeVar +from typing_extensions import NotRequired, Required + +class Unrelated: pass + +class Open(TypedDict): + x: Required[int] + y: Required[str] + z: NotRequired[str] + +T = TypeVar('T', bound=Open) + +def f(d: T, direct: Open, u: Unrelated) -> None: + x: Literal['x'] + y: Literal['y'] + z: Literal['z'] + missing: Literal['missing'] + x_or_y: Literal['x', 'y'] + x_or_z: Literal['x', 'z'] + y_or_missing: Literal['y', 'missing'] + s = '' + + reveal_type(d.get(x)) # N: Revealed type is "builtins.int" + reveal_type(direct.get(x)) # N: Revealed type is "builtins.int" + reveal_type(d.get(y, u)) # N: Revealed type is "builtins.str" + reveal_type(direct.get(y, u)) # N: Revealed type is "builtins.str" + reveal_type(d.get(z)) # N: Revealed type is "builtins.str | None" + reveal_type(direct.get(z)) # N: Revealed type is "builtins.str | None" + reveal_type(d.get(z, u)) # N: Revealed type is "builtins.str | __main__.Unrelated" + reveal_type(direct.get(z, u)) # N: Revealed type is "builtins.str | __main__.Unrelated" + reveal_type(d.get(z, 'a')) # N: Revealed type is "builtins.str" + reveal_type(direct.get(z, 'a')) # N: Revealed type is "builtins.str" + reveal_type(d.get(x_or_y)) # N: Revealed type is "builtins.int | builtins.str" + reveal_type(direct.get(x_or_y)) # N: Revealed type is "builtins.int | builtins.str" + reveal_type(d.get(x_or_y, u)) # N: Revealed type is "builtins.int | builtins.str" + reveal_type(direct.get(x_or_y, u)) # N: Revealed type is "builtins.int | builtins.str" + reveal_type(d.get(x_or_z)) # N: Revealed type is "builtins.int | builtins.str | None" + reveal_type(direct.get(x_or_z)) # N: Revealed type is "builtins.int | builtins.str | None" + reveal_type(d.get(x_or_z, u)) # N: Revealed type is "builtins.int | builtins.str | __main__.Unrelated" + reveal_type(direct.get(x_or_z, u)) # N: Revealed type is "builtins.int | builtins.str | __main__.Unrelated" + reveal_type(d.get(missing)) # N: Revealed type is "builtins.object" + reveal_type(direct.get(missing)) # N: Revealed type is "builtins.object" + reveal_type(d.get('missing')) # N: Revealed type is "builtins.object" + reveal_type(direct.get('missing')) # N: Revealed type is "builtins.object" + reveal_type(d.get(y_or_missing)) # N: Revealed type is "builtins.object" + reveal_type(direct.get(y_or_missing)) # N: Revealed type is "builtins.object" + reveal_type(d.get(y_or_missing, u)) # N: Revealed type is "builtins.object" + reveal_type(direct.get(y_or_missing, u)) # N: Revealed type is "builtins.object" + reveal_type(d.get(s)) # N: Revealed type is "builtins.object" + reveal_type(direct.get(s)) # N: Revealed type is "builtins.object" +[builtins fixtures/dict.pyi] +[typing fixtures/typing-typeddict.pyi] + +[case testTypedDictGetBoundTypeVarInvalidArgs] +from typing import TypedDict, TypeVar + +class MyDict(TypedDict): + x: int + y: str + +T = TypeVar('T', bound=MyDict) + +def f(d: T) -> None: + d.get() # E: All overload variants of "get" of "Mapping" require at least one argument \ + # N: Possible overload variants: \ + # N: def get(self, 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(self, str, /) -> object \ + # N: def get(self, str, object, /) -> object \ + # N: def [V] get(self, str, int | V, /) -> object + d.get('x', 1 + 'y') # E: Unsupported operand types for + ("int" and "str") + s = '' + reveal_type(d.get(s)) # N: Revealed type is "builtins.object" + reveal_type(d.get('z')) # N: Revealed type is "builtins.object" +[builtins fixtures/dict.pyi] +[typing fixtures/typing-typeddict.pyi] + +[case testTypedDictGetBoundTypeVarUnrelated] +from typing import TypeVar + +D = TypeVar('D', bound=dict[str, int]) + +def f(d: D) -> None: + reveal_type(d.get('x')) # N: Revealed type is "builtins.int | None" + reveal_type(d.get('x', 1)) # N: Revealed type is "builtins.int" +[builtins fixtures/dict.pyi] +[typing fixtures/typing-typeddict.pyi]