Reputation: 2511
I am looking to make a Python (3.8+) function that is:
TypedDict
Here is a code sample to help explain:
from typing import Any, Literal, TypedDict
class Foo(TypedDict):
bar: int
baz: str
spam: Any
foo = Foo(bar=0, baz="hi", spam=1.0)
def get_type(key: Literal["bar", "baz"]): # How to type hint the return here?
"""Function that get TypedDict's value when passed a key."""
val = foo[key]
# This works via intelligent indexing
# SEE: https://mypy.readthedocs.io/en/stable/literal_types.html#intelligent-indexing
reveal_type(val) # mypy: Revealed type is 'Union[builtins.int, builtins.str]'
return val
fetched_type = get_type("bar")
reveal_type(fetched_type) # mypy: Revealed type is 'Any'
# I would like this to have output: 'int'
In case you can't tell, the static type checker I am using is mypy
.
The function I have above get_type
gets halfway there with intelligent indexing, but I am not sure how to type hint the return of get_type
.
What type hint should I put for the return of get_type
?
Research
Both of these questions
provide answers using TypeVar
. Is there some way of using TypeVar
with TypedDict
?
Upvotes: 4
Views: 2038
Reputation: 32303
If I understood your problem correctly, you could use @overload
for get_type
:
from typing import Any, Literal, TypedDict, Union, overload
class Foo(TypedDict):
bar: int
baz: str
spam: Any
foo = Foo(bar=0, baz="hi", spam=1.0)
@overload
def get_type(key: Literal["bar"]) -> int: ...
@overload
def get_type(key: Literal["baz"]) -> str: ...
def get_type(key: Literal["bar", "baz"]) -> Union[int, str]:
"""Function that get TypedDict's value when passed a key."""
val = foo[key]
reveal_type(val) # mypy: Revealed type is 'Union[builtins.int, builtins.str]'
return val
fetched_type = get_type("bar")
reveal_type(fetched_type) # mypy: Revealed type is 'builtins.int'
Upvotes: 2