如何为含自定义键/值的字典加类型提示？兼容Python3.9与检查工具

针对需求——在Python 3.9下实现支持枚举、指定字面量字符串作为键（或混合），同时兼容mypy、pylint且无需用户手动标注字典类型——可以通过重载函数签名+运行时校验的方式实现，既保证向后兼容性，又维持类型安全。

核心思路 #

1. 定义枚举类型和对应字面量集合，覆盖新旧两种键类型
2. 利用@overload提供多套函数签名，匹配不同的字典键类型场景
3. 实现函数时使用宽泛的类型接受输入，同时添加运行时校验确保键的合法性
4. 支持值类型为str或int的组合场景

完整代码实现 #

from typing import Union, Dict, Literal, overload
from enum import Enum

# 新数据类型：枚举
class Animals(Enum):
    cat = "cat"
    dog = "dog"
    snake = "snake"

# 兼容旧版本：允许的字面量字符串键
AnimalsLiteral = Literal["cat", "dog", "snake"]

# 组合键类型与值类型
AnimalKey = Union[Animals, AnimalsLiteral]
AnimalValue = Union[str, int]

# 重载函数签名，匹配不同场景
@overload
def dummy(a: Dict[Animals, AnimalValue]) -> None: ...

@overload
def dummy(a: Dict[AnimalsLiteral, AnimalValue]) -> None: ...

@overload
def dummy(a: Dict[str, AnimalValue]) -> None: ...

@overload
def dummy(a: Dict[AnimalKey, AnimalValue]) -> None: ...

# 实际实现函数，添加运行时校验
def dummy(a: Dict[Union[Animals, str], AnimalValue]) -> None:
    # 预定义允许的字符串键集合
    allowed_str_keys = {member.value for member in Animals}
    
    # 运行时校验所有键的合法性
    for key in a.keys():
        if isinstance(key, str):
            if key not in allowed_str_keys:
                raise ValueError(f"非法字符串键：{key}，仅允许 {allowed_str_keys}")
        elif not isinstance(key, Animals):
            raise TypeError(f"非法键类型：{type(key)}，仅允许 Animals 枚举或指定字符串")

# 测试用例
if __name__ == "__main__":
    # 1. 键为Animals枚举的字典（值为str）
    input_data1 = {Animals.dog: "dog"}
    dummy(input_data1)
    
    # 2. 键为字面量字符串的字典（值为int）
    input_data2 = {"dog": 123}
    dummy(input_data2)
    
    # 3. 混合键：枚举+字面量字符串（值混合str/int）
    input_data3 = {Animals.dog: "dog", "snake": 456}
    dummy(input_data3)
    
    # 4. 非法键（运行时报错）
    # input_data4 = {"bird": "bird"}
    # dummy(input_data4)

关键说明 #

1. 重载签名的作用：

   • 覆盖了纯枚举键、纯字面量键、普通字符串键（兼容用户未标注类型的场景）、混合键四种情况
   • mypy会自动匹配最贴合的签名，避免对用户传入的{"dog": "dog"}这类未标注类型的字典报错

2. 运行时校验的必要性：

   • 因为我们允许接受Dict[str, AnimalValue]类型，需要在运行时确保字符串键属于允许的字面量集合，避免非法输入
   • 枚举键的类型已经由Python本身保证，无需额外校验

3. 兼容性：

   • 完全符合Python 3.9要求，未使用StrEnum或TypeAlias
   • 兼容mypy、pylint的类型检查规则
   • 用户无需手动标注字典类型，直接传入普通字典即可通过类型检查

扩展：更严格的类型约束（可选） #

如果希望在类型层面完全限制字符串键为指定字面量，同时避免用户手动标注，可以通过mypy配置文件添加--allow-subtyped-unions选项，然后将函数签名简化为：

def dummy(a: Union[Dict[Animals, AnimalValue], Dict[AnimalsLiteral, AnimalValue]]) -> None:
    # 运行时校验逻辑不变
    ...

此配置会让mypy允许将dict[str, str]兼容到Dict[AnimalsLiteral, str]，但需要团队统一mypy配置。

内容的提问来源于stack exchange，提问作者Figue