Typing a Python dictionary as Dict[key[T], value[K]] where T and K are restricted

I have a dictionary that keeps track of two different types of values (let's say city and village) using two different types of keys (city_key and village_key). I'd like to annotate this dictionary with generics, so that when the dictionary receives key of type city_key mypy/Pyright should annotate the return value as city. Likewise, if you try to assign a city value to a village_key, mypy/pyright should throw an error.

The alternative is to maintain two different dictionaries, one for cities and one for villages, but I am curious if I can get away with one dictionary.

There is a question just like mine here, but it went unanswered.

Some pseudo code to show what I am aiming for in practice:

# two types of aliased keys
# ... edited to use NewType as per juanpa.arrivillaga comment

CityKey = NewType("CityKey", str)
VillageKey = NewType("VillageKey", str)

# two types of values, city and village
class City:...
class Village:...

# key generator that returns city or village key based on type of input
def generate_key(settlement: City | Village) -> CityKey | VillageKey: ...

# declare some keys & values
london = City("London")
london_key = generate_key(london)
mousehole = Village("Mousehole")
mousehole_key = generate_key(village)

# instantiate the dictionary
data: [????] = {}

# assign city to city key, and village to village key
data[london_key] = london
data[mousehole_key] = mousehole

# trying to assign village to city key should raise a type check error
data[london_key] = mousehole

# type of value accessed by village key should be village
reveal_type(data[mousehole_key]) # Type[Village]

Solution

You can use typing.overload for this purpose, which can help us go from types like Callable[[A1 | B1], A2 | B2] to one where it can be either Callable[[A1], A2] or Callable[[B1], B2], and a subclass of dict.

from typing import overload

@overload
def generate_key(settlement: City) -> CityKey:
    # Just a stub
    ...


@overload
def generate_key(settlement: Village) -> VillageKey:
    # Just a stub
    ...


def generate_key(settlement):
    # Contains the actual implementation
    [...]


class CityOrVillageDict(dict):
    @overload
    def __setitem__(self, key: CityKey, value: City) -> None:
        # Just a stub
        ...

    @overload
    def __setitem__(self, key: VillageKey, value: Village) -> None:
        # Just a stub
        ...

    def __setitem__(self, key, value):
        # Overloaded functions need an implementation
        super().__setitem__(key, value)

    @overload
    def __getitem__(self, key: CityKey) -> City:
        # Just a stub
        ...

    @overload
    def __getitem__(self, key: VillageKey) -> Village:
        # Just a stub
        ...

    def __getitem__(self, key):
        # Overloaded functions need an implementation
        return super().__getitem__(key)

data = CityOrVillageDict()