Conditional types with mypy

Question

I have the following code snippet:

from typing import TypedDict

class Super(TypedDict):
    foo: int

class SubA(Super):
    bar: int

class SubB(Super):
    zap: int

def print_props(inp: Super, key: str):
    print(inp[key])

When I call the method print_props with either an instance of SubA or SubB it would be valid as they are sub types of Super.

But mypy will complain about inp[key] as the key must be literal "foo". Is it possible to give mypy hints so that it is capable of deciding which keys are valid? For example: "When print_props is called with an instance of SubB only "foo" and "zap" are valid."

I took a look at generics; I think it is possible to declare a type variable that is restricted to sub types of Super, but is it possible to express the dependency between the concrete type of the type variable (SubA or SubB) and the literal values key should then be restricted to?

Answer 1

Overloads with Literal could well do it, though I do wonder if a different design would be better. I'm a bit concerned about the increasingly frequent usage of overload and Literal in SO answers. They both suggest a design smell to me

@overload
def printMyProps(input: SubA, key: Literal["foo", "bar"]) -> None: ...

@overload
def printMyProps(input: SubB, key: Literal["foo", "zap"]) -> None: ...

def printMyProps(input: SubA | SubB, key: Literal["foo", "bar", "zap"]) -> None:
  print(input[key])  # type: ignore

I've used type: ignore because it's a short function and I can't use isinstance on TypedDict. TBH overload implementations often require type hacks. The API works as intended though