Python 以枚举为键的字典

Question

假设我有一个枚举

class Color(Enum):
    RED = "RED"
    GREEN = "GREEN"
    BLUE = "BLUE"

我想创建一个 ColorDict class 作为本机 python 字典，但只将 Color 枚举或其对应的字符串值作为键。

d = ColorDict() # I want to implement a ColorDict class such that ...

d[Color.RED] = 123
d["RED"] = 456  # I want this to override the previous value
d[Color.RED]    # ==> 456
d["foo"] = 789  # I want this to produce an KeyError exception

实现此 ColorDict class 的“pythonic 方式”是什么？我应该使用继承（覆盖 python 的原生 dict）还是组合（保留 dict 作为成员）？

Answer 1

一种方法是使用抽象基class collections.abc.MutableMapping，这样，您只需要覆盖抽象方法，然后您就可以确保访问始终通过您的逻辑——您也可以使用 dict 执行此操作，但是例如，覆盖 dict.__setitem__ 将 而不是 影响 dict.update、dict.setdefault 等...所以你也必须手动覆盖那些。通常，只使用抽象基 class:

更容易

from collections.abc import MutableMapping
from enum import Enum

class Color(Enum):
    RED = "RED"
    GREEN = "GREEN"
    BLUE = "BLUE"

class ColorDict(MutableMapping):

    def __init__(self): # could handle more ways of initializing  but for simplicity...
        self._data = {}

    def __getitem__(self, item):
        return self._data[color]

    def __setitem__(self, item, value):
        color = self._handle_item(item)
        self._data[color] = value

    def __delitem__(self, item):
        del self._data[color]

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

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

    def _handle_item(self, item):
        try:
            color = Color(item)
        except ValueError:
            raise KeyError(item) from None
        return color

注意，您还可以添加：

    def __repr__(self):
        return repr(self._data)

为了更容易调试。

回复中的一个例子：

In [3]: d = ColorDict() # I want to implement a ColorDict class such that ...
   ...:
   ...: d[Color.RED] = 123
   ...: d["RED"] = 456  # I want this to override the previous value
   ...: d[Color.RED]    # ==> 456
Out[3]: 456

In [4]: d["foo"] = 789  # I want this to produce an KeyError exception
   ...:
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-4-9cf80d6dd8b4> in <module>
----> 1 d["foo"] = 789  # I want this to produce an KeyError exception

<ipython-input-2-a0780e16594b> in __setitem__(self, item, value)
     17
     18     def __setitem__(self, item, value):
---> 19         color = self._handle_item(item)
     20         self._data[color] = value
     21

<ipython-input-2-a0780e16594b> in _handle_item(self, item)
     34             color = Color(item)
     35         except ValueError:
---> 36             raise KeyError(item) from None
     37         return color
     38     def __repr__(self): return repr(self._data)

KeyError: 'foo'

In [5]: d
Out[5]: {<Color.RED: 'RED'>: 456}

Answer 2

一个简单的解决方案是稍微修改您的 Color 对象，然后 subclass dict 添加密钥测试。我会做这样的事情：

class Color(Enum):
    RED = "RED"
    GREEN = "GREEN"
    BLUE = "BLUE"

    @classmethod
    def is_color(cls, color):
        if isinstance(color, cls):
            color=color.value
        if not color in cls.__members__:
            return False
        else:
            return True


class ColorDict(dict):
    
    def __setitem__(self, k, v):
        if Color.is_color(k):
            super().__setitem__(Color(k), v)
        else:
            raise KeyError(f"Color {k} is not valid")

    def __getitem__(self, k):
        if isinstance(k, str):
            k = Color(k.upper())
        return super().__getitem__(k)

d = ColorDict()

d[Color.RED] = 123
d["RED"] = 456
d[Color.RED]
d["foo"] = 789

在Colorclass中，我添加了一个测试函数到returnTrue或False如果一个颜色is/isn' t 在允许列表中。 upper() 函数将字符串设为大写，以便与预定义值进行比较。

然后我有 subclassed dict 对象来覆盖 __setitem__ 特殊方法以包括对传递值的测试，以及 __getitem__ 的覆盖将作为 str 传递的任何密钥转换为正确的 Enum。根据您希望如何使用 ColorDict class 的具体细节，您可能需要覆盖更多函数。这里有一个很好的解释：How to properly subclass dict and override __getitem__ & __setitem__

Python 以枚举为键的字典

Python dictionary with enum as key

python

enums

dictionary