将 Python 字典键入为 Dict[key[T], value[K]] ，其中 T 和 K 受到限制

Question

问题是：

我有一个字典，它使用两种不同类型的键（city_key 和 village_key).我想用泛型注释这本字典，这样当字典接收到 city_key 类型的键时，mypy/Pyright 应该将 return 值注释为 city。同样，如果您尝试将 city 值分配给 village_key，mypy/pyright 应该会引发错误。

另一种方法是维护两本不同的字典，一本用于城市，一本用于乡村，但我很好奇我是否可以只用一本字典。

到目前为止我尝试了什么

我研究了这个问题，但找不到明确的例子/答案。有一个问题和我的一样 here，但没有人回答。

一些伪代码来展示我在实践中的目标


# 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]

谢谢！

PS 我不确定我是否在标题中正确使用了术语 covariant。如果我不是，请告诉我。

Answer 1

您可以为此目的使用 typing.overload，这可以帮助我们从 Callable[[A1 | B1], A2 | B2] 等类型转变为可以是 Callable[[A1], A2] 或 Callable[[B1], B2] 的类型，并且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()

将 Python 字典键入为 Dict[key[T], value[K]] ，其中 T 和 K 受到限制

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

python

mypy

python-typing

pyright