Python - TypedDict type hint keys

I have 2 classes. One that inherits TypedDict class in which I have set of keys that should accept certain type of data:

class UserType(TypeDict):
   first_name: str
   last_name: str
   age: int

In another class Customer, I am using UserType dict for storing the information of the customer and have a method to update the information:

class Customer:
   def __init__():
     self._userInfo: UserType = {
       "first_name": "John",
       "last_name": "John",
       "age": 27
    }
    
    def set_info(key: str, val: Any):
      self._userInfo[key] = val

In set_info I want to type hint key to show only keys that are defined in UserType like:

customer = Customer()
customer.set_info(SHOULD HINT KEY HERE (first_name, ...), "Somedata")

dict. I tried something like:

def set_info(key: UserType, val: Any):
   ...

But this didn't work. I guess, I am trying to use it as interfaces as I am coming from strict typed language background. Any idea how this can be achieved? I am flexible to change it something different other than a dict

Solution

TypedDict doesn't expose its keys as a distinct type, nor is there really a good way to define such a type.

Not recommended

Some people use typing.Literal, but that is intended for, well, literals. Use key: Literal["first_name", "last_name", "age"] allows c.set_info("first_name", "bob"), but not k = "first_name"; c.set_info(k, "bob"), because k is not a literal. I don't recommend doing this, but I felt it needed to be addressed.

Better

Another alternative would be to explicitly define a separate enumerate type and use its members, rather than raw keys, as arguments.

class CustomerKey(enum.StrEnum):
    FIRST_NAME = "first_name"
    LAST_NAME = "last_name"
    AGE = "age"

# This could (somewhat awkwardly) also be defined as
#
#   CustomerKey = enum.StrEnum(dict(zip(map(str.upper, UserType.__annotation__),
#                                       UserType.__annotation__)))
#
# to avoid repetition of the exact keys.


class Customer:
    ...

    def set_info(key: CustomerKey, val: Any):     
        self._userInfo[key] = val

Better still

This is assuming you don't have a good reason to store the attributes separately, and only use the TypedDict as the argument to a class method, for example,

@dataclasses.dataclass
class Customer:
    first_name: str
    last_name: str
    age: int

    @classmethod
    def from_dict(cls, d: UserType):
        return cls(**d)

Best (hypothetical)

Don't read too much into what follows. I'm mostly thinking out loud about ways that TypeDict could expose the information it contains as type hints.

In an ideal world, the keys of a TypedDict would themselves form a type. Then perhaps you could write something like

class Customer:
    ...

    def set_info(key: UserType.keys, val: Any):
        ...

Even better, there would also be a way to tied the value assigned to key to a type that could be used to annotate val, so that if key == "first_name", then val must have type `str. I imagine this would require some mix of dependent types (where types can depend on value), singleton types (where for every value there is an associated type that contains only that value) and generics. Something vaguely like

    def set_info(key: UserType.keys, val: UserType[key])