Python如何实现仅接受指定Enum枚举及对应字符串为键的自定义字典？

方案选择 #

优先选择基于collections.abc.MutableMapping的组合模式实现，而非直接继承原生dict：

• 原生dict的大量内置方法（如update、setdefault等）由C底层实现，不会调用用户重写的__setitem__，仅重写核心方法会出现校验逻辑遗漏，导致非法键被写入
• 继承MutableMapping只需要实现5个基础抽象方法，所有dict标准行为会自动适配自定义的键校验逻辑，不会出现功能遗漏，代码更简洁可维护

完整实现代码 #

from enum import Enum
from collections.abc import MutableMapping

# 你定义的Color枚举
class Color(Enum):
    RED = "RED"
    GREEN = "GREEN"
    BLUE = "BLUE"

class ColorDict(MutableMapping):
    def __init__(self, *args, **kwargs):
        # 内部用普通dict做存储
        self._storage = dict()
        # 兼容dict的初始化传参逻辑
        self.update(*args, **kwargs)

    def _validate_transform_key(self, key):
        """校验键合法性，统一转换为枚举对应的字符串值"""
        # 键是Color枚举成员，直接取对应值
        if isinstance(key, Color):
            return key.value
        # 键是字符串，校验是否是合法的枚举名称
        if isinstance(key, str) and key in Color.__members__:
            return key
        # 非法键抛出KeyError
        raise KeyError(f"Invalid key {key!r}: only Color enum members or their matching string values are allowed")

    def __getitem__(self, key):
        return self._storage[self._validate_transform_key(key)]

    def __setitem__(self, key, value):
        self._storage[self._validate_transform_key(key)] = value

    def __delitem__(self, key):
        del self._storage[self._validate_transform_key(key)]

    def __iter__(self):
        return iter(self._storage)

    def __len__(self):
        return len(self._storage)

    # 可选：自定义输出格式，和普通dict表现一致
    def __repr__(self):
        return f"{type(self).__name__}({repr(self._storage)})"

测试验证 #

运行你给出的用例，完全符合预期：

d = ColorDict()

d[Color.RED] = 123
d["RED"] = 456
print(d[Color.RED])  # 输出 456

try:
    d["foo"] = 789
except KeyError as e:
    print(e)  # 抛出KeyError: Invalid key 'foo': ...

另外所有dict原生方法都可以正常使用，且自动携带键校验：

# 正常执行
d.update({"GREEN": 789, Color.BLUE: 1011})
print(d.get(Color.GREEN))  # 输出789

# 会抛出KeyError
d.update({"invalid_key": 123})

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