如何使用附加方法从现有字典创建 Python Enum 类?

was*_*zil 6 python enums mixins

假设我有一个预先存在的映射作为字典:

value_map = {'a': 1, 'b': 2}
Run Code Online (Sandbox Code Playgroud)

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

from enum import Enum
MyEnum = Enum('MyEnum', value_map)
Run Code Online (Sandbox Code Playgroud)

并像这样使用它

a = MyEnum.a
print(a.value)
>>> 1
print(a.name)
>>> 'a'
Run Code Online (Sandbox Code Playgroud)

但后来我想为我的新枚举类定义一些方法:

def double_value(self):
    return self.value * 2
Run Code Online (Sandbox Code Playgroud)

当然,我可以这样做:

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

但正如我所说,我必须使用预定义的值映射字典,所以我不能这样做。如何才能实现这一目标?我尝试从另一个像 mixin 一样定义此方法的类继承,但我无法弄清楚。

Mar*_*ers 8

您可以将具有 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
Run Code Online (Sandbox Code Playgroud)

对于从不使用语句的全功能方法,您可以使用以下函数class创建基本混合:type()

DoubledEnum = type('DoubledEnum', (), {'double_value': property(double_value)})
MyEnum = enum.Enum('MyEnum', value_map, type=DoubledEnum)
Run Code Online (Sandbox Code Playgroud)

您还可以以enum.EnumMeta()相同的方式使用元类,就像 Python 创建class MyEnum(enum.Enum): ...子类时的方式一样:

  1. 使用元类挂钩创建类字典__prepare__
  2. 调用元类,传入类名、基类((enum.Enum,)此处)以及在步骤 1 中创建的类字典。

使用的自定义字典子类enum.EnumMeta并不是真正为易于重用而设计的;它实现了一个__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)
Run Code Online (Sandbox Code Playgroud)

然后传递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
Run Code Online (Sandbox Code Playgroud)

否则,您可以创建不带成员的枚举的子类(任何描述符都不是成员,因此函数、属性、类方法等),因此您可以首先定义不带成员的枚举:

class DoubledEnum(enum.Enum):
    @property
    def double_value(self):
        return self.value * 2
Run Code Online (Sandbox Code Playgroud)

对于功能 API(例如enum.Enum(..., type=DoubledEnum))和我编码为 的元类方法来说,这是可接受的基类enum_with_extras()。