Is it possible to type annotate Python function parameter used as TypedDict key to make mypy happy?

Question

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

Answer 1

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 (playground):

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]]"