While working through code challenges I am trying to use type annotations for all function parameters/return types. I use mypy in strict mode with the goal of no errors.
I've spent some time on this one and can't figure it out - example of problem:
from typing import Literal, NotRequired, TypedDict
class Movie(TypedDict):
Title: str
Runtime: str
Awards: str
class MovieData(TypedDict):
Awards: NotRequired[str]
Runtime: NotRequired[str]
def get_movie_field(movies: list[Movie], field: Literal['Awards', 'Runtime']) -> dict[str, MovieData]:
return {movie['Title']: {field: movie[field]} for movie in movies}
# Check with mypy 1.14.0:
PS> mypy --strict .\program.py
program.py:13: error: Expected TypedDict key to be string literal [misc]
Found 1 error in 1 file (checked 1 source file)
Is there some way to make this work? I realize I can append # type: ignore
to the line but I'm curious if there's another way.
I've tried a bunch of different permutations such as ensuring field matches one of the literal values before the dict comprehension but nothing seems to help.
I like @chepner's solution. Unfortunately, I couldn't figure out how to make it work with mypy in strict mode - taking his example:
def get_movie_field(movies, field):
return {movie['Title']: {field: movie[field]} for movie in movies}
I get:
PS> mypy --strict program.py
program.py:18: error: Function is missing a type annotation [no-untyped-def]
Found 1 error in 1 file (checked 1 source file)
As the last get_movie_field
isn't type annotated.
From reviewing the mypy docs, I updated the last get_movie_field
function as follows - but it still doesn't fix the problem:
# @chepner's full solution updated with type annotations for
# last get_movie_field:
from typing import Literal, NotRequired, TypedDict, overload
class Movie(TypedDict):
Title: str
Runtime: str
Awards: str
class MovieData(TypedDict):
Awards: NotRequired[str]
Runtime: NotRequired[str]
@overload
def get_movie_field(movies: list[Movie], field: Literal['Awards']) -> dict[str, MovieData]:
...
@overload
def get_movie_field(movies: list[Movie], field: Literal['Runtime']) -> dict[str, MovieData]:
...
def get_movie_field(movies: list[Movie], field: Literal['Awards']|Literal['Runtime']
) -> dict[str, MovieData]:
return {movie['Title']: {field: movie[field]} for movie in movies}
PS> mypy --strict program.py
program.py:19: error: Expected TypedDict key to be string literal [misc]
Found 1 error in 1 file (checked 1 source file)
However, this inspired me to find another approach. I'm not saying it's great but it does work in the sense that mypy doesn't complain:
from typing import Literal, NotRequired, TypedDict
class Movie(TypedDict):
Title: str
Runtime: str
Awards: str
class MovieData(TypedDict):
Awards: NotRequired[str]
Runtime: NotRequired[str]
def create_MovieData(key: str, value: str) -> MovieData:
"""Creates a new TypedDict using a given key and value."""
if key == 'Awards':
return {'Awards': value}
elif key == 'Runtime':
return {'Runtime': value}
raise ValueError(f"Invalid key: {key}")
def get_movie_field(movies: list[Movie], field: Literal['Awards', 'Runtime']) -> dict[str, MovieData]:
return {movie['Title']: create_MovieData(field, movie[field]) for movie in movies}
PS> mypy --strict program.py
Success: no issues found in 1 source file
You can use TypeVar
to help mypy deduce the exact literal that is being used, and this can even be propagated to the output. That way mypy knows which keys are definitely present in the output dicts and which are not.
Instead of:
class MovieData(TypedDict):
Awards: NotRequired[str]
Runtime: NotRequired[str]
You would use:
K = TypeVar('K', Literal['Runtime'], Literal['Awards'])
# NB. NOT TypeVar('K', Literal['Runtime', 'Awards']) -- TypeError
# AND NOT TypeVar('K', bound=Literal['Runtime', 'Awards']) -- coalesces to str
MovieData = dict[K, str]
Full example:
from typing import TypedDict, TypeVar, Literal
class Movie(TypedDict):
Title: str
Runtime: str
Awards: str
K = TypeVar('K', Literal['Runtime'], Literal['Awards'])
MovieData = dict[K, str]
def extract(movies: list[Movie], key: K) -> list[MovieData[K]]:
return [{key: movie[key]} for movie in movies]
res = extract([], key='Runtime')
reveal_type(res)
# note: Revealed type is "list[dict[Literal['Runtime'], str]]"