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

Question

我有一个字典，它使用两种不同类型的键（

city

和

village

）来跟踪两种不同类型的值（比如说

city_key

和

village_key

）。我想用泛型注释这个字典，这样当字典接收到

city_key

类型的键时 mypy/Pyright 应该将返回值注释为

city

。同样，如果您尝试将

city

值分配给

village_key

，mypy/pyright 应该会抛出错误。

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

有一个和我一样的问题这里，但没有得到解答。

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

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

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()