输入一个Python字典，如Dict[key[T]， value[K]]，其中T和K是受限的

Question

问题是:

我有一个字典，它使用两种不同类型的键(city_key和village_key)跟踪两种不同类型的值(假设city和village)。我想用泛型注释这个字典，这样当字典接收到类型为city_key的键时，myypy/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]

谢谢!

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