invert-typed-dict.py

from collections.abc import Sequence
from typing import Protocol, Sequence, TypedDict, TypeVar, cast


class LiteralProtocol(Protocol):
    __args__: Sequence[str]


class TypedDictProtocol(Protocol):
    __annotations__: dict[str, LiteralProtocol]


_T = TypeVar('_T', bound=TypedDictProtocol)


def invert_dict_with_literals(td: type[_T], indices: Sequence[int] | int = 0) -> _T:
    # Still must use cast(). Fundamentally it is because the bound type on _T cannot have generics.
    if isinstance(indices, Sequence):
        return cast(_T, {
            k: v.__args__[i]
            for i, (k, v) in zip(indices, td.__annotations__.items(), strict=True)
        })
    return cast(_T, {k: v.__args__[indices] for k, v in td.__annotations__.items()})
  
  class ATitleTitleDescDict(TypedDict):
      title: Literal['a title']
      description: Literal['some description']
        
reveal_type(invert_dict_with_literals(ATitleTitleDescDict)) # ATitleTitleDescDict