Python: Annotate variable as key of a TypedDict

Question

Basically a distilled down version of this (as yet unanswered) question.

I want to state that a variable should only take on values that are keys in a TypedDict.

At present I'm defining a separate Literal type to represent the keys, for example:

from typing import Literal, TypedDict


class MyTD(TypedDict):
    a: int
    b: int


mytd = MyTD(a=1, b=2)

key = "a"

mytd[key]  # error: TypedDict key must be a string literal; expected one of ('a', 'b')

MyTDKeyT = Literal["a", "b"]

typed_key: MyTDKeyT = "b"

mytd[typed_key]  # no error

I would like to be able to replace the Literal definition for all the usual reasons of wanting to minimize duplicated code.

Pseudo-code:

key: Keys[MyTD] = "a"
mytd[key]  # would be no error
not_key: Keys[MyTD] = "z"  # error

Is there a way to achieve this?

To clarify, given that mypy can tell me that the key type needs to be a literal of "a" or "b", I'm hoping there might be a less error prone way to annotate a variable to that type, rather than having to maintain two separate lists of keys side-by-side, once in the TypedDict definition, once in the Literal definition.

Answer 1

Using MyPy, I don't think this is possible. I ran this experiment:

from typing import TypedDict

class MyTD(TypedDict):
    a: str
    b: int

d = MyTD(a='x', b=2)
reveal_type(list(d))

The MyPy output was:

Revealed type is "builtins.list[builtins.str]"

This indicates that internally it is not tracking the keys as literals. Otherwise, we would expect:

Revealed type is "builtins.list[Literal['A', 'B']]" 

Also, this errors out in MyPy, so __required_keys__ isn't even inspectable:

reveal_type(MyTD.__required_keys__)