CODEWITHSUNDEEP
pythondictionarymypypython-typing
Intrastellar Explorer
Intrastellar Explorer

Reputation: 2511

Function for getting type of TypedDict value that works with static type checker

I am looking to make a Python (3.8+) function that is:

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

Answers (1)

alex_noname
alex_noname

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

Related Questions