如何使用其他方法从现有字典创建 Python 枚举 class?

How to create Python Enum class from existing dict with additional methods?

比方说,我有一个预先存在的映射作为字典:

value_map = {'a': 1, 'b': 2}

我可以这样创建一个枚举 class:

from enum import Enum
MyEnum = Enum('MyEnum', value_map)

并像这样使用它

a = MyEnum.a
print(a.value)
>>> 1
print(a.name)
>>> 'a'

但是我想为我的新枚举定义一些方法class:

def double_value(self):
    return self.value * 2

我当然可以:

class MyEnum(Enum):
    a = 1
    b = 2
    @property
    def double_value(self):
        return self.value * 2

但是正如我所说,我必须使用预定义的值映射字典,所以我不能这样做。 如何实现?我试图从另一个 class 继承这个方法,像混合一样定义这个方法,但我想不通。

您可以使用 mixin 方法将基类型传入函数 API,使用 type 参数:

>>> import enum
>>> value_map = {'a': 1, 'b': 2}
>>> class DoubledEnum:
...     @property
...     def double_value(self):
...         return self.value * 2
...
>>> MyEnum = enum.Enum('MyEnum', value_map, type=DoubledEnum)
>>> MyEnum.a.double_value
2

对于从不使用 class 语句的全功能方法,您可以使用 type() function:

创建基本混合
DoubledEnum = type('DoubledEnum', (), {'double_value': property(double_value)})
MyEnum = enum.Enum('MyEnum', value_map, type=DoubledEnum)

您也可以使用 enum.EnumMeta() metaclass 以相同的方式使用 Python 创建 class MyEnum(enum.Enum): ... subclass 时的方式:

  1. 使用元class __prepare__ hook
  2. 创建class字典
  3. 调用 metaclass,传入 class 名称、碱基(此处为 (enum.Enum,))和在步骤 1 中创建的 class 字典。

enum.EnumMeta 使用的自定义字典 ​​subclass 并不是真正为便于重用而设计的;它实现了一个 __setitem__ 钩子来记录元数据,但没有覆盖 dict.update() 方法,所以我们在使用你的 value_map 字典时需要小心:

import enum

def enum_with_extras(name, value_map, bases=enum.Enum, **extras):
    if not isinstance(bases, tuple):
        bases = bases,
    if not any(issubclass(b, enum.Enum) for b in bases):
        bases += enum.Enum,
    classdict = enum.EnumMeta.__prepare__(name, bases)
    for key, value in {**value_map, **extras}.items():
        classdict[key] = value
    return enum.EnumMeta(name, bases, classdict)

然后将 double_value=property(double_value) 传递给该函数(连同枚举名称和 value_map 字典):

>>> def double_value(self):
...     return self.value * 2
...
>>> MyEnum = enum_with_extras('MyEnum', value_map, double_value=property(double_value))
>>> MyEnum.a
<MyEnum.a: 1>
>>> MyEnum.a.double_value
2

否则您可以创建没有成员的枚举的子class(任何 descriptor 都不是成员,因此函数、属性、class 方法等.), 所以你可以先定义一个没有成员的枚举:

class DoubledEnum(enum.Enum):
    @property
    def double_value(self):
        return self.value * 2

对于函数 API(例如 enum.Enum(..., type=DoubledEnum))和我编码为 [=29= 的 metaclass 方法,这是可接受的基础 class ].

您可以创建一个新的元 class(使用元元class 或工厂函数,就像我在下面做的那样)派生自 enum.EnumMeta(元class 用于枚举)并在创建 class

之前添加成员
import enum
import collections.abc


def enum_metaclass_with_default(default_members):
    """Creates an Enum metaclass where `default_members` are added"""
    if not isinstance(default_members, collections.abc.Mapping):
        default_members = enum.Enum('', default_members).__members__

    default_members = dict(default_members)

    class EnumMetaWithDefaults(enum.EnumMeta):
        def __new__(mcs, name, bases, classdict):
            """Updates classdict adding the default members and
            creates a new Enum class with these members
            """

            # Update the classdict with default_members
            # if they don't already exist
            for k, v in default_members.items():
                if k not in classdict:
                    classdict[k] = v

            # Add `enum.Enum` as a base class

            # Can't use `enum.Enum` in `bases`, because
            # that uses `==` instead of `is`
            bases = tuple(bases)
            for base in bases:
                if base is enum.Enum:
                    break
            else:
                bases = (enum.Enum,) + bases

            return super(EnumMetaWithDefaults, mcs).__new__(mcs, name, bases, classdict)

    return EnumMetaWithDefaults


value_map = {'a': 1, 'b': 2}


class MyEnum(metaclass=enum_metaclass_with_default(value_map)):
    @property
    def double_value(self):
        return self.value * 2


assert MyEnum.a.double_value == 2

另一种解决方案是直接尝试并更新 locals(),因为当您尝试分配值时,它会被创建枚举值的映射所取代。

import enum


value_map = {'a': 1, 'b': 2}


def set_enum_values(locals, value_map):
    # Note that we can't use `locals.update(value_map)`
    # because it's `locals.__setitem__(k, v)` that
    # creates the enum value, and `update` doesn't
    # call `__setitem__`.
    for k, v in value_map:
        locals[k] = v


class MyEnum(enum.Enum):
    set_enum_values(locals(), value_map)

    @property
    def double_value(self):
        return self.value * 2


assert MyEnum.a.double_value == 2

这似乎定义得很好,a = 1 很可能与 locals()['a'] = 1 相同,但将来可能会改变。第一个解决方案更健壮且更简单(我还没有在其他 Python 实现中测试过它,但它可能工作相同)

PLUS: 在@Artyer 的回答中添加更多内容(粗略的技巧)。

请注意,如果您从 dict 创建 Enum,您还可以为它提供“附加”功能,请参阅...

from enum import Enum

_colors = {"RED": (1, "It's the color of blood."), "BLUE": (2, "It's the color of the sky.")}


def _set_members_colors(locals: dict):
    for k, v in colors.items():
        locals[k] = v[0]


class Colors(int, Enum):

    _set_members_colors(locals())

    @property
    def description(self):
        return colors[self.name][1]

print(str(Colors.RED))
print(str(Colors.RED.value))
print(str(Colors.RED.description))

输出...

Colors.RED
1
It's the color of blood.

谢谢!