如何用mypy静态检查字典包含Union类型所有选项的键？

我定义了由多个dataclass组成的Union类型PossibleOption，想要创建以这些类的实例为键的字典，需要用mypy静态检查该字典是否包含Union中每个选项对应的键。现有代码如下：

from dataclasses import dataclass
from typing import Union

@dataclass
class Option1:
    def __hash__(self):   # 不用太关注这个方法，字典需要可哈希的键，仅此而已
        return hash(type(self))

@dataclass
class Option2:
    def __hash__(self):
        return hash(type(self))

@dataclass
class Option3:
    def __hash__(self):
        return hash(type(self))

PossibleOption = Union[Option1, Option2, Option3]

# 我想要创建这样的字典：
a = {Option1(): "a", Option2(): "bb", Option3(): "cc"}
# 比如 print(a[Option2()]) 会输出 "bb"

我知道可以通过定义接收PossibleOption为参数的函数，使用match语句返回对应值，mypy能检查是否覆盖所有选项：

def equivalent_to_dictionary(po: PossibleOption) -> str:
    match po:
        case Option1():
            return "aa"
        case Option2():
            return "bb"
        case Option3():
            return "cc"

但我希望保持用户使用方式更简洁，请问是否有办法用mypy静态检查字典的键是否包含Union类型的每个选项？

可行的解决方案 #

方法1：构建函数+match语句强制覆盖 #

编写一个专门的字典构建函数，利用mypy对match语句的覆盖检查特性，间接保证字典包含所有必要键。如果Union类型新增了选项，mypy会在match处提示未覆盖的case，从而避免遗漏键：

from dataclasses import dataclass
from typing import Union, Dict

@dataclass
class Option1:
    def __hash__(self):
        return hash(type(self))

@dataclass
class Option2:
    def __hash__(self):
        return hash(type(self))

@dataclass
class Option3:
    def __hash__(self):
        return hash(type(self))

PossibleOption = Union[Option1, Option2, Option3]

def build_option_dict() -> Dict[PossibleOption, str]:
    option_dict: Dict[PossibleOption, str] = {}
    # 遍历所有Option类，创建实例并添加到字典
    for option_cls in [Option1, Option2, Option3]:
        instance = option_cls()
        match instance:
            case Option1():
                option_dict[instance] = "a"
            case Option2():
                option_dict[instance] = "bb"
            case Option3():
                option_dict[instance] = "cc"
    return option_dict

# 使用时直接调用函数获取字典
a = build_option_dict()

方法2：TypedDict+类型推导转换 #

由于你的Option实例哈希仅依赖类型，可先通过TypedDict强制类类型的键完整，再转换成实例为键的字典，兼顾类型检查和使用简洁性：

from dataclasses import dataclass
from typing import Union, Literal, TypedDict

@dataclass
class Option1:
    def __hash__(self):
        return hash(type(self))

@dataclass
class Option2:
    def __hash__(self):
        return hash(type(self))

@dataclass
class Option3:
    def __hash__(self):
        return hash(type(self))

# 定义类类型的Literal
OptionType = Literal[Option1, Option2, Option3]
# 定义total=True的TypedDict，强制包含所有类类型的键
class OptionTypeDict(TypedDict, total=True):
    Option1: str
    Option2: str
    Option3: str

# 先构建符合类型约束的类-值映射
type_map: OptionTypeDict = {
    Option1: "a",
    Option2: "bb",
    Option3: "cc"
}

# 转换成实例为键的目标字典
a = {cls(): val for cls, val in type_map.items()}

这种方式下，若遗漏某个类的键，mypy会直接报错，同时最终的字典使用方式和你原本的需求完全一致。

方法3：自定义mypy插件（进阶） #

如果上述方法仍不符合需求，可编写自定义mypy插件，实现对字典键是否覆盖Union所有选项的直接检查。不过该方案需要编写插件代码，学习成本较高，适合复杂场景。

内容的提问来源于stack exchange，提问作者Pierre-olivier Gendraud