如何扩展 Python Enum?

在 Python 3.4中扩展 Enum类型的最佳实践是什么,是否存在这样做的可能性?

例如:

from enum import Enum


class EventStatus(Enum):
success = 0
failure = 1


class BookingStatus(EventStatus):
duplicate = 2
unknown = 3


Traceback (most recent call last):
...
TypeError: Cannot extend enumerations

目前还没有可能的方法来创建一个包含成员的基础枚举类并在其他枚举类中使用它(如上面的示例所示)。还有其他实现 Python 枚举继承的方法吗?

67212 次浏览

Subclassing an enumeration is allowed only if the enumeration does not define any members.

Allowing subclassing of enums that define members would lead to a violation of some important invariants of types and instances.

https://docs.python.org/3/library/enum.html#restricted-enum-subclassing

So no, it's not directly possible.

While uncommon, it is sometimes useful to create an enum from many modules. The aenum1 library supports this with an extend_enum function:

from aenum import Enum, extend_enum


class Index(Enum):
DeviceType    = 0x1000
ErrorRegister = 0x1001


for name, value in (
('ControlWord', 0x6040),
('StatusWord', 0x6041),
('OperationMode', 0x6060),
):
extend_enum(Index, name, value)


assert len(Index) == 5
assert list(Index) == [Index.DeviceType, Index.ErrorRegister, Index.ControlWord, Index.StatusWord, Index.OperationMode]
assert Index.DeviceType.value == 0x1000
assert Index.StatusWord.value == 0x6041

1 Disclosure: I am the author of the Python stdlib Enum, the enum34 backport, and the Advanced Enumeration (aenum) library.

Calling the Enum class directly and making use of chain allows the extension (joining) of an existing enum.

I came upon the problem of extending enums while working on a CANopen implementation. Parameter indices in the range from 0x1000 to 0x2000 are generic to all CANopen nodes while e.g. the range from 0x6000 onwards depends open whether the node is a drive, io-module, etc.

nodes.py:

from enum import IntEnum


class IndexGeneric(IntEnum):
""" This enum holds the index value of genric object entrys
"""
DeviceType    = 0x1000
ErrorRegister = 0x1001


Idx = IndexGeneric

drives.py:

from itertools import chain
from enum import IntEnum
from nodes import IndexGeneric


class IndexDrives(IntEnum):
""" This enum holds the index value of drive object entrys
"""
ControlWord   = 0x6040
StatusWord    = 0x6041
OperationMode = 0x6060


Idx= IntEnum('Idx', [(i.name, i.value) for i in chain(IndexGeneric,IndexDrives)])

I've opted to use a metaclass approach to this problem.

from enum import EnumMeta


class MetaClsEnumJoin(EnumMeta):
"""
Metaclass that creates a new `enum.Enum` from multiple existing Enums.


@code
from enum import Enum


ENUMA = Enum('ENUMA', {'a': 1, 'b': 2})
ENUMB = Enum('ENUMB', {'c': 3, 'd': 4})
class ENUMJOINED(Enum, metaclass=MetaClsEnumJoin, enums=(ENUMA, ENUMB)):
pass


print(ENUMJOINED.a)
print(ENUMJOINED.b)
print(ENUMJOINED.c)
print(ENUMJOINED.d)
@endcode
"""


@classmethod
def __prepare__(metacls, name, bases, enums=None, **kargs):
"""
Generates the class's namespace.
@param enums Iterable of `enum.Enum` classes to include in the new class.  Conflicts will
be resolved by overriding existing values defined by Enums earlier in the iterable with
values defined by Enums later in the iterable.
"""
#kargs = {"myArg1": 1, "myArg2": 2}
if enums is None:
raise ValueError('Class keyword argument `enums` must be defined to use this metaclass.')
ret = super().__prepare__(name, bases, **kargs)
for enm in enums:
for item in enm:
ret[item.name] = item.value  #Throws `TypeError` if conflict.
return ret


def __new__(metacls, name, bases, namespace, **kargs):
return super().__new__(metacls, name, bases, namespace)
#DO NOT send "**kargs" to "type.__new__".  It won't catch them and
#you'll get a "TypeError: type() takes 1 or 3 arguments" exception.


def __init__(cls, name, bases, namespace, **kargs):
super().__init__(name, bases, namespace)
#DO NOT send "**kargs" to "type.__init__" in Python 3.5 and older.  You'll get a
#"TypeError: type.__init__() takes no keyword arguments" exception.

This metaclass can be used like so:

>>> from enum import Enum
>>>
>>> ENUMA = Enum('ENUMA', {'a': 1, 'b': 2})
>>> ENUMB = Enum('ENUMB', {'c': 3, 'd': 4})
>>> class ENUMJOINED(Enum, metaclass=MetaClsEnumJoin, enums=(ENUMA, ENUMB)):
...     e = 5
...     f = 6
...
>>> print(repr(ENUMJOINED.a))
<ENUMJOINED.a: 1>
>>> print(repr(ENUMJOINED.b))
<ENUMJOINED.b: 2>
>>> print(repr(ENUMJOINED.c))
<ENUMJOINED.c: 3>
>>> print(repr(ENUMJOINED.d))
<ENUMJOINED.d: 4>
>>> print(repr(ENUMJOINED.e))
<ENUMJOINED.e: 5>
>>> print(repr(ENUMJOINED.f))
<ENUMJOINED.f: 6>

This approach creates a new Enum using the same name-value pairs as the source Enums, but the resulting Enum members are still unique. The names and values will be the same, but they will fail direct comparisons to their origins following the spirit of Python's Enum class design:

>>> ENUMA.b.name == ENUMJOINED.b.name
True
>>> ENUMA.b.value == ENUMJOINED.b.value
True
>>> ENUMA.b == ENUMJOINED.b
False
>>> ENUMA.b is ENUMJOINED.b
False
>>>

Note what happens in the event of a namespace conflict:

>>> ENUMC = Enum('ENUMA', {'a': 1, 'b': 2})
>>> ENUMD = Enum('ENUMB', {'a': 3})
>>> class ENUMJOINEDCONFLICT(Enum, metaclass=MetaClsEnumJoin, enums=(ENUMC, ENUMD)):
...     pass
...
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
File "<stdin>", line 19, in __prepare__
File "C:\Users\jcrwfrd\AppData\Local\Programs\Python\Python37\lib\enum.py", line 100, in __setitem__
raise TypeError('Attempted to reuse key: %r' % key)
TypeError: Attempted to reuse key: 'a'
>>>

This is due to the base enum.EnumMeta.__prepare__ returning a special enum._EnumDict instead of the typical dict object that behaves different upon key assignment. You may wish to suppress this error message by surrounding it with a try-except TypeError, or there may be a way to modify the namespace before calling super().__prepare__(...).

I think you could do it in this way:

from typing import List
from enum import Enum


def extend_enum(current_enum, names: List[str], values: List = None):
if not values:
values = names


for item in current_enum:
names.append(item.name)
values.append(item.value)


return Enum(current_enum.__name__, dict(zip(names, values)))


class EventStatus(Enum):
success = 0
failure = 1


class BookingStatus(object):
duplicate = 2
unknown = 3


BookingStatus = extend_enum(EventStatus, ['duplicate','unknown'],[2,3])

the key points is:

  • python could change anything at runtime
  • class is object too

Another way :

Letter = Enum(value="Letter", names={"A": 0, "B": 1})
LetterExtended = Enum(value="Letter", names=dict({"C": 2, "D": 3}, **{i.name: i.value for i in Letter}))

Or :

LetterDict = {"A": 0, "B": 1}
Letter = Enum(value="Letter", names=LetterDict)


LetterExtendedDict = dict({"C": 2, "D": 3}, **LetterDict)
LetterExtended = Enum(value="Letter", names=LetterExtendedDict)

Output :

>>> Letter.A
<Letter.A: 0>
>>> Letter.C
Traceback (most recent call last):
File "<input>", line 1, in <module>
File "D:\jhpx\AppData\Local\Programs\Python\Python36\lib\enum.py", line 324, in __getattr__
raise AttributeError(name) from None
AttributeError: C
>>> LetterExtended.A
<Letter.A: 0>
>>> LetterExtended.C
<Letter.C: 2>

I tested that way on 3.8. We may inherit existing enum but we need to do it also from base class (at last position).

Docs:

A new Enum class must have one base Enum class, up to one concrete data type, and as many object-based mixin classes as needed. The order of these base classes is:

class EnumName([mix-in, ...,] [data-type,] base-enum):
pass

Example:

class Cats(Enum):
SIBERIAN = "siberian"
SPHINX = "sphinx"




class Animals(Cats, Enum):
LABRADOR = "labrador"
CORGI = "corgi"

After that you may access Cats from Animals:

>>> Animals.SIBERIAN
<Cats.SIBERIAN: 'siberian'>

But if you want to iterate over this enum, only new members were accessible:

>>> list(Animals)
[<Animals.LABRADOR: 'labrador'>, <Animals.CORGI: 'corgi'>]

Actually this way is for inheriting methods from base class, but you may use it for members with these restrictions.

Another way (a bit hacky)

As described above, to write some function to join two enums in one. I've wrote that example:

def extend_enum(inherited_enum):
def wrapper(added_enum):
joined = {}
for item in inherited_enum:
joined[item.name] = item.value
for item in added_enum:
joined[item.name] = item.value
return Enum(added_enum.__name__, joined)
return wrapper




class Cats(Enum):
SIBERIAN = "siberian"
SPHINX = "sphinx"




@extend_enum(Cats)
class Animals(Enum):
LABRADOR = "labrador"
CORGI = "corgi"

But here we meet another problems. If we want to compare members it fails:

>>> Animals.SIBERIAN == Cats.SIBERIAN
False

Here we may compare only names and values of newly created members:

>>> Animals.SIBERIAN.value == Cats.SIBERIAN.value
True

But if we need iteration over new Enum, it works ok:

>>> list(Animals)
[<Animals.SIBERIAN: 'siberian'>, <Animals.SPHINX: 'sphinx'>, <Animals.LABRADOR: 'labrador'>, <Animals.CORGI: 'corgi'>]

So choose your way: simple inheritance, inheritance emulation with decorator (recreation in fact), or adding a new dependency like aenum (I haven't tested it, but I expect it support all features I described).

For correct type specification, you could use the Union operator:

from enum import Enum
from typing import Union


class EventStatus(Enum):
success = 0
failure = 1


class BookingSpecificStatus(Enum):
duplicate = 2
unknown = 3


BookingStatus = Union[EventStatus, BookingSpecificStatus]


example_status: BookingStatus
example_status = BookingSpecificStatus.duplicate
example_status = EventStatus.success

Yes, you can modify an Enum. The example code, below, is somewhat hacky and it obviously depends on internals of Enum which it has no business whatsoever to depend on. On the other hand, it works.

class ExtIntEnum(IntEnum):
@classmethod
def _add(cls, value, name):
obj = int.__new__(cls, value)
obj._value_ = value
obj._name_ = name
obj.__objclass__ = cls


cls._member_map_[name] = obj
cls._value2member_map_[value] = obj
cls._member_names_.append(name)


class Fubar(ExtIntEnum):
foo = 1
bar = 2


Fubar._add(3,"baz")
Fubar._add(4,"quux")

Specifically, observe the obj = int.__new__() line. The enum module jumps through a few hoops to find the correct __new__ method for the class that should be enumerated. We ignore these hoops here because we already know how integers (or rather, instances of subclasses of int) are created.

It's a good idea not to use this in production code. If you have to, you really should add guards against duplicate values or names.

I wanted to inherit from Django's IntegerChoices which is not possible due to the "Cannot extend enumerations" limitation. I figured it could be done by a relative simple metaclass.

CustomMetaEnum.py:

class CustomMetaEnum(type):
def __new__(self, name, bases, namespace):
# Create empty dict to hold constants (ex. A = 1)
fields = {}


# Copy constants from the namespace to the fields dict.
fields = {key:value for key, value in namespace.items() if isinstance(value, int)}
    

# In case we're about to create a subclass, copy all constants from the base classes' _fields.
for base in bases:
fields.update(base._fields)


# Save constants as _fields in the new class' namespace.
namespace['_fields'] = fields
return super().__new__(self, name, bases, namespace)


# The choices property is often used in Django.
# If other methods such as values(), labels() etc. are needed
# they can be implemented below (for inspiration [Django IntegerChoice source][1])
@property
def choices(self):
return [(value,key) for key,value in self._fields.items()]

main.py:

from CustomMetaEnum import CustomMetaEnum


class States(metaclass=CustomMetaEnum):
A = 1
B = 2
C = 3


print("States: ")
print(States.A)
print(States.B)
print(States.C)
print(States.choices)




print("MoreStates: ")
class MoreStates(States):
D = 22
pass


print(MoreStates.A)
print(MoreStates.B)
print(MoreStates.C)
print(MoreStates.D)
print(MoreStates.choices)

python3.8 main.py:

States:
1
2
3
[(1, 'A'), (2, 'B'), (3, 'C')]
MoreStates:
1
2
3
22
[(22, 'D'), (1, 'A'), (2, 'B'), (3, 'C')]

You can't extend enums but you can create a new one by merging them.
Tested for Python 3.6

from enum import Enum




class DummyEnum(Enum):
a = 1




class AnotherDummyEnum(Enum):
b = 2




def merge_enums(class_name: str, enum1, enum2, result_type=Enum):
if not (issubclass(enum1, Enum) and issubclass(enum2, Enum)):
raise TypeError(
f'{enum1} and {enum2} must be derived from Enum class'
)


attrs = {attr.name: attr.value for attr in set(chain(enum1, enum2))}
return result_type(class_name, attrs, module=__name__)




result_enum = merge_enums(
class_name='DummyResultEnum',
enum1=DummyEnum,
enum2=AnotherDummyEnum,
)

Plenty of good answers here already but here's another one purely using Enum's Functional API.

Probably not the most beautiful solution but it avoids code duplication, works out of the box, no additional packages/libraries are need, and it should be sufficient to cover most use cases:

from enum import Enum


class EventStatus(Enum):
success = 0
failure = 1


BookingStatus = Enum(
"BookingStatus",
[es.name for es in EventStatus] + ["duplicate", "unknown"],
start=0,
)


for bs in BookingStatus:
print(bs.name, bs.value)


# success 0
# failure 1
# duplicate 2
# unknown 3

If you'd like to be explicit about the values assigned, you can use:

BookingStatus = Enum(
"BookingStatus",
[(es.name, es.value) for es in EventStatus] + [("duplicate", 6), ("unknown", 7)],
)


for bs in BookingStatus:
print(bs.name, bs.value)


# success 0
# failure 1
# duplicate 6
# unknown 7

Decorator to extend Enum

To expand on Mikhail Bulygin's answer, a decorator can be used to extend an Enum (and support equality by using a custom Enum base class).

1. Enum base class with value-based equality

from enum import Enum
from typing import Any




class EnumBase(Enum):
def __eq__(self, other: Any) -> bool:
if isinstance(other, Enum):
return self.value == other.value
return False

2. Decorator to extend Enum class

from typing import Callable


def extend_enum(parent_enum: EnumBase) -> Callable[[EnumBase], EnumBase]:
"""Decorator function that extends an enum class with values from another enum class."""
def wrapper(extended_enum: EnumBase) -> EnumBase:
joined = {}
for item in parent_enum:
joined[item.name] = item.value
for item in extended_enum:
joined[item.name] = item.value
return EnumBase(extended_enum.__name__, joined)
return wrapper

Example

>>> from enum import Enum
>>> from typing import Any, Callable
>>> class EnumBase(Enum):
def __eq__(self, other: Any) -> bool:
if isinstance(other, Enum):
return self.value == other.value
return False
>>> def extend_enum(parent_enum: EnumBase) -> Callable[[EnumBase], EnumBase]:
def wrapper(extended_enum: EnumBase) -> EnumBase:
joined = {}
for item in parent_enum:
joined[item.name] = item.value
for item in extended_enum:
joined[item.name] = item.value
return EnumBase(extended_enum.__name__, joined)
return wrapper
>>> class Parent(EnumBase):
A = 1
B = 2
>>> @extend_enum(Parent)
class ExtendedEnum(EnumBase):
C = 3
>>> Parent.A == ExtendedEnum.A
True
>>> list(ExtendedEnum)
[<ExtendedEnum.A: 1>, <ExtendedEnum.B: 2>, <ExtendedEnum.C: 3>]