Is there a way to do the same code below using Dataclass instead of Enum?
from enum import Enum, auto
class State(Enum):
val_A= auto()
val_B = auto()
val_C = auto()
The only solution I found is the following code:
from dataclasses import dataclass
#dataclass(frozen=True)
class State():
val_A:str = 'val_A'
val_B:str = 'val_B'
val_C:str = 'val_C'
thank you for the suggestions.
Descriptors
One approach could be to use a descriptor class, defined as below:
class Auto:
_GLOBAL_STATE = {}
__slots__ = ('_private_name', )
# `owner` is the class or type, whereas instance is an object of `owner`
def __get__(self, instance, owner):
try:
return getattr(instance, self._private_name)
except AttributeError:
_state = self.__class__._GLOBAL_STATE
_dflt = _state[owner] = _state.get(owner, 0) + 1
return _dflt
def __set_name__(self, owner, name):
self._private_name = '_' + name
def __set__(self, instance, value):
# use object.__setattr__() instead of setattr() as dataclasses
# also does, in case of a "frozen" dataclass
object.__setattr__(instance, self._private_name, value)
Usage would be as follows:
from dataclasses import dataclass
#dataclass(frozen=True)
class State:
val_A: int = Auto()
val_B: int = Auto()
val_C: int = Auto()
s = State()
print(s) # State(val_A=1, val_B=2, val_C=3)
assert s.val_B == 2
s = State(val_A=5)
assert s.val_A == 5
assert s.val_C == 3
Optimal Approach
The most performant approach I can think of, would be to replace the default values before dataclasses is able to process the class.
Initially this is O(N) time, as it would require iterating over all the class members (including dataclass fields) at least once. However, the real benefit is that it replaces the default values for auto values, such as val_A: int = 1, before the #dataclass decorator is able to process the class.
For example, define a metaclass such as one below:
# sentinel value to detect when to replace a field's default
auto = object()
def check_auto(name, bases, cls_dict):
default = 1
for name, val in cls_dict.items():
if val == auto:
cls_dict[name] = default
default += 1
cls = type(name, bases, cls_dict)
return cls
Usage is as below:
from dataclasses import dataclass
#dataclass(frozen=True)
class State(metaclass=check_auto):
val_A: int = auto
val_B: int = auto
val_C: int = auto
s = State()
print(s) # State(val_A=1, val_B=2, val_C=3)
assert s.val_B == 2
s = State(val_A=5)
assert s.val_A == 5
assert s.val_C == 3
I want to know a simple way to make a dataclass bar frozen.
#dataclass
class Bar:
foo: int
bar = Bar(foo=1)
In other words, I want a function like the following some_fn_to_freeze
frozen_bar = some_fn_to_freeze(bar)
frozen_bar.foo = 2 # Error
And, an inverse function some_fn_to_unfreeze
bar = som_fn_to_unfrozen(frozen_bar)
bar.foo = 3 # not Error
The standard way to mutate a frozen dataclass is to use dataclasses.replace:
old_bar = Bar(foo=123)
new_bar = dataclasses.replace(old_bar, foo=456)
assert new_bar.foo == 456
For more complex use-cases, you can use the dataclass utils module from: https://github.com/google/etils
It add a my_dataclass = my_dataclass.unfrozen() member, which allow to mutate frozen dataclasses directly
# pip install etils[edc]
from etils import edc
#edc.dataclass(allow_unfrozen=True) # Add the `unfrozen()`/`frozen` method
#dataclasses.dataclass(frozen=True)
class A:
x: Any = None
y: Any = None
old_a = A(x=A(x=A()))
# After a is unfrozen, the updates on nested attributes will be propagated
# to the top-level parent.
a = old_a.unfrozen()
a.x.x.x = 123
a.x.y = 'abc'
a = a.frozen() # `frozen()` recursively call `dataclasses.replace`
# Only the `unfrozen` object is mutated. Not the original one.
assert a == A(x=A(x=A(x = 123), y='abc'))
assert old_a == A(x=A(x=A()))
As seen in the example, you can return unfrozen/frozen copies of the dataclass, which was explicitly designed to mutate nested dataclasses.
#edc.dataclass also add a a.replace(**kwargs) method to the dataclass (alias of dataclasses.dataclass)
a = A()
a = a.replace(x=123, y=456)
assert a == A(x=123, y=456)
dataclass doesn't have built-in support for that. Frozen-ness is tracked on a class-wide basis, not per-instance, and there's no support for automatically generating frozen or unfrozen equivalents of dataclasses.
While you could try to do something to generate new dataclasses on the fly, it'd interact very poorly with isinstance, ==, and other things you'd want to work. It's probably safer to just write two dataclasses and converter methods:
#dataclass
class Bar:
foo: int
def as_frozen(self):
return FrozenBar(self.foo)
#dataclass(frozen=True)
class FrozenBar:
foo: int
def as_unfrozen(self):
return Bar(self.foo)
Python dataclasses are great, but the attrs package is a more flexible alternative, if you are able to use a third-party library. For example:
import attr
# Your class of interest.
#attr.s()
class Bar(object):
val = attr.ib()
# A frozen variant of it.
#attr.s(frozen = True)
class FrozenBar(Bar):
pass
# Three instances:
# - Bar.
# - FrozenBar based on that Bar.
# - Bar based on that FrozenBar.
b1 = Bar(123)
fb = FrozenBar(**attr.asdict(b1))
b2 = Bar(**attr.asdict(fb))
# We can modify the Bar instances.
b1.val = 777
b2.val = 888
# Check current vals.
for x in (b1, fb, b2):
print(x)
# But we cannot modify the FrozenBar instance.
try:
fb.val = 999
except attr.exceptions.FrozenInstanceError:
print(fb, 'unchanged')
Output:
Bar(val=888)
FrozenBar(val=123)
Bar(val=999)
FrozenBar(val=123) unchanged
I'm using the following code to get a frozen copy of a dataclass class or instance:
import dataclasses
from dataclasses import dataclass, fields, asdict
import typing
from typing import TypeVar
FDC_SELF = TypeVar('FDC_SELF', bound='FreezableDataClass')
#dataclass
class FreezableDataClass:
#classmethod
def get_frozen_dataclass(cls: Type[FDC_SELF]) -> Type[FDC_SELF]:
"""
#return: a generated frozen dataclass definition, compatible with the calling class
"""
cls_fields = fields(cls)
frozen_cls_name = 'Frozen' + cls.__name__
frozen_dc_namespace = {
'__name__': frozen_cls_name,
'__module__': __name__,
}
excluded_from_freezing = cls.attrs_excluded_from_freezing()
for attr in dir(cls):
if attr.startswith('__') or attr in excluded_from_freezing:
continue
attr_def = getattr(cls, attr)
if hasattr(attr_def, '__func__'):
attr_def = classmethod(getattr(attr_def, '__func__'))
frozen_dc_namespace[attr] = attr_def
frozen_dc = dataclasses.make_dataclass(
cls_name=frozen_cls_name,
fields=[(f.name, f.type, f) for f in cls_fields],
bases=(),
namespace=frozen_dc_namespace,
frozen=True,
)
globals()[frozen_dc.__name__] = frozen_dc
return frozen_dc
#classmethod
def attrs_excluded_from_freezing(cls) -> typing.Iterable[str]:
return tuple()
def get_frozen_instance(self: FDC_SELF) -> FDC_SELF:
"""
#return: an instance of a generated frozen dataclass, compatible with the current dataclass, with copied values
"""
cls = type(self)
frozen_dc = cls.get_frozen_dataclass()
# noinspection PyArgumentList
return frozen_dc(**asdict(self))
Derived classes could overwrite attrs_excluded_from_freezing to exclude methods which wouldn't work on a frozen dataclass.
Why didn't I prefer other existing answers?
3rd party libraries - etils.edc, If I would use a solution from one of the previous answers, it would be this one. E.g. to get the ability to recursively freeze/unfreeze.
3rd party libraries - attrs
duplicated code
In this article Nick Coghlan talks about some of the design decisions that went in to the PEP 435 Enum type, and how EnumMeta can be subclassed to provide a different Enum experience.
However, the advice I give (and I am the primary stdlib Enum author) about using a metaclass is it should not be done without a really good reason -- such as not being able to accomplish what you need with a class decorator, or a dedicated function to hide any ugliness; and in my own work I've been able to do whatever I needed simply by using __new__, __init__, and/or normal class/instance methods when creating the Enum class:
Enum with attributes
Handling missing members
class constants that are not Enum members
And then there is this cautionary tale of being careful when delving into Enum, with and without metaclass subclassing:
Is it possible to override __new__ in an enum to parse strings to an instance?
Given all that, when would I need to fiddle with EnumMeta itself?
The best cases I have seen so far for subclassing EnumMeta comes from these four questions:
A more pythonic way to define an enum with dynamic members
Prevent invalid enum attribute assignment
Create an abstract Enum class
Invoke a function when an enum member is accessed
We'll examine the dynamic member case further here.
First, a look at the code needed when not subclassing EnumMeta:
The stdlib way
from enum import Enum
import json
class BaseCountry(Enum):
def __new__(cls, record):
member = object.__new__(cls)
member.country_name = record['name']
member.code = int(record['country-code'])
member.abbr = record['alpha-2']
member._value_ = member.abbr, member.code, member.country_name
if not hasattr(cls, '_choices'):
cls._choices = {}
cls._choices[member.code] = member.country_name
cls._choices[member.abbr] = member.country_name
return member
def __str__(self):
return self.country_name
Country = BaseCountry(
'Country',
[(rec['alpha-2'], rec) for rec in json.load(open('slim-2.json'))],
)
The aenum way 1 2
from aenum import Enum, MultiValue
import json
class Country(Enum, init='abbr code country_name', settings=MultiValue):
_ignore_ = 'country this' # do not add these names as members
# create members
this = vars()
for country in json.load(open('slim-2.json')):
this[country['alpha-2']] = (
country['alpha-2'],
int(country['country-code']),
country['name'],
)
# have str() print just the country name
def __str__(self):
return self.country_name
The above code is fine for a one-off enumeration -- but what if creating Enums from JSON files was common for you? Imagine if you could do this instead:
class Country(JSONEnum):
_init_ = 'abbr code country_name' # remove if not using aenum
_file = 'some_file.json'
_name = 'alpha-2'
_value = {
1: ('alpha-2', None),
2: ('country-code', lambda c: int(c)),
3: ('name', None),
}
As you can see:
_file is the name of the json file to use
_name is the path to whatever should be used for the name
_value is a dictionary mapping paths to values3
_init_ specifies the attribute names for the different value components (if using aenum)
The JSON data is taken from https://github.com/lukes/ISO-3166-Countries-with-Regional-Codes -- here is a short excerpt:
[{"name":"Afghanistan","alpha-2":"AF","country-code":"004"},
{"name":"Åland Islands","alpha-2":"AX","country-code":"248"},
{"name":"Albania","alpha-2":"AL","country-code":"008"},
{"name":"Algeria","alpha-2":"DZ","country-code":"012"}]
Here is the JSONEnumMeta class:
class JSONEnumMeta(EnumMeta):
#classmethod
def __prepare__(metacls, cls, bases, **kwds):
# return a standard dictionary for the initial processing
return {}
def __init__(cls, *args , **kwds):
super(JSONEnumMeta, cls).__init__(*args)
def __new__(metacls, cls, bases, clsdict, **kwds):
import json
members = []
missing = [
name
for name in ('_file', '_name', '_value')
if name not in clsdict
]
if len(missing) in (1, 2):
# all three must be present or absent
raise TypeError('missing required settings: %r' % (missing, ))
if not missing:
# process
name_spec = clsdict.pop('_name')
if not isinstance(name_spec, (tuple, list)):
name_spec = (name_spec, )
value_spec = clsdict.pop('_value')
file = clsdict.pop('_file')
with open(file) as f:
json_data = json.load(f)
for data in json_data:
values = []
name = data[name_spec[0]]
for piece in name_spec[1:]:
name = name[piece]
for order, (value_path, func) in sorted(value_spec.items()):
if not isinstance(value_path, (list, tuple)):
value_path = (value_path, )
value = data[value_path[0]]
for piece in value_path[1:]:
value = value[piece]
if func is not None:
value = func(value)
values.append(value)
values = tuple(values)
members.append(
(name, values)
)
# get the real EnumDict
enum_dict = super(JSONEnumMeta, metacls).__prepare__(cls, bases, **kwds)
# transfer the original dict content, _items first
items = list(clsdict.items())
items.sort(key=lambda p: (0 if p[0][0] == '_' else 1, p))
for name, value in items:
enum_dict[name] = value
# add the members
for name, value in members:
enum_dict[name] = value
return super(JSONEnumMeta, metacls).__new__(metacls, cls, bases, enum_dict, **kwds)
# for use with both Python 2/3
JSONEnum = JSONEnumMeta('JsonEnum', (Enum, ), {})
A few notes:
JSONEnumMeta.__prepare__ returns a normal dict
EnumMeta.__prepare__ is used to get an instance of _EnumDict -- this is the proper way to get one
keys with a leading underscore are passed to the real _EnumDict first as they may be needed when processing the enum members
Enum members are in the same order as they were in the file
1 Disclosure: I am the author of the Python stdlib Enum, the enum34 backport, and the Advanced Enumeration (aenum) library.
2 This requires aenum 2.0.5+.
3 The keys are numeric to keep multiple values in order should your Enum need more than one.
Is it possible to extend classes created using the new Enum functionality in Python 3.4? How?
Simple subclassing doesn't appear to work. An example like
from enum import Enum
class EventStatus(Enum):
success = 0
failure = 1
class BookingStatus(EventStatus):
duplicate = 2
unknown = 3
will give an exception like TypeError: Cannot extend enumerations or (in more recent versions) TypeError: BookingStatus: cannot extend enumeration 'EventStatus'.
How can I make it so that BookingStatus reuses the enumeration values from EventStatus and adds more?
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/howto/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 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
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
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__(...).
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 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
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,
)
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>]
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')]
Conceptually, it does not make sense to extend an enumeration in this sense. The problem is that this violates the Liskov Substitution Principle: instances of a subclass are supposed to be usable anywhere an instance of the base class could be used, but an instance of BookingStatus could not reliably be used anywhere that an EventStatus is expected. After all, if that instance had a value of BookingStatus.duplicate or BookingStatus.unknown, that would not be a valid enumeration value for an EventStatus.
We can create a new class that reuses the EventStatus setup by using the functional API. A basic example:
event_status_codes = [s.name for s in EventStatus]
BookingStatus = Enum(
'BookingStatus', event_status_codes + ['duplicate', 'unknown']
)
This approach re-numbers the enumeration values, ignoring what they were in EventStatus. We can also pass name-value pairs in order to specify the enum values; this lets us do a bit more analysis, in order to reuse the old values and auto-number new ones:
def extend_enum(result_name, base, *new_names):
base_values = [(v.name, v.value) for v in base]
next_number = max(v.value for v in base) + 1
new_values = [(name, i) for i, name in enumerate(new_names, next_number)]
return Enum(result_name, base_values + new_values)
# Now we can do:
BookingStatus = extend_enum('BookingStatus', EventStatus, 'duplicate', 'unknown')
I've been messing around with python's enum library and have come across a conundrum. In the docs, they show an example of an auto-numbering enum, wherein something is defined:
class Color(AutoNumber):
red = ()
green = ()
...
I want to make a similar class, but the value would automatically be set from the name of the member AND keep the functionality that you get from doing the str and enum mixin stuff
So something like:
class Animal(MagicStrEnum):
horse = ()
dog = ()
Animal.dog == 'dog' # True
I've looked at the source code of the enum module and tried a lot of variations messing around with __new__ and the EnumMeta class
Update: 2017-03-01
In Python 3.6 (and Aenum 2.01) Flag and IntFlag classes have been added; part of that was a new auto() helper that makes this trivially easy:
>>> class AutoName(Enum):
... def _generate_next_value_(name, start, count, last_values):
... return name
...
>>> class Ordinal(AutoName):
... NORTH = auto()
... SOUTH = auto()
... EAST = auto()
... WEST = auto()
...
>>> list(Ordinal)
[<Ordinal.NORTH: 'NORTH'>, <Ordinal.SOUTH: 'SOUTH'>, <Ordinal.EAST: 'EAST'>, <Ordinal.WEST: 'WEST'>]
Original answer
The difficulty with an AutoStr class is that the name of the enum member is not passed into the code that creates it, so it is unavailable for use. Another wrinkle is that str is immutable, so we can't change those types of enums after they have been created (by using a class decorator, for example).
The easiest thing to do is use the Functional API:
Animal = Enum('Animal', [(a, a) for a in ('horse', 'dog')], type=str)
which gives us:
>>> list(Animal)
[<Animal.horse: 'horse'>, <Animal.dog: 'dog'>]
>>> Animal.dog == 'dog'
True
The next easiest thing to do, assuming you want to make a base class for your future enumeration use, would be something like my DocEnem:
class DocEnum(Enum):
"""
compares equal to all cased versions of its name
accepts a doctring for each member
"""
def __new__(cls, *args):
"""Ignores arguments (will be handled in __init__)"""
obj = object.__new__(cls)
obj._value_ = None
return obj
def __init__(self, doc=None):
# first, fix _value_
self._value_ = self._name_.lower()
self.__doc__ = doc
def __eq__(self, other):
if isinstance(other, basestring):
return self._value_ == other.lower()
elif not isinstance(other, self.__class__):
return NotImplemented
return self is other
def __hash__(self):
# keep DocEnum hashable
return hash(self._value_)
def __ne__(self, other):
return not self == other
and in use:
class SpecKind(DocEnum):
REQUIRED = "required value"
OPTION = "single value per name"
MULTI = "multiple values per name (list form)"
FLAG = "boolean value per name"
KEYWORD = 'unknown options'
Note that unlike the first option, DocEnum members are not strs.
If you want to do it the hard way: subclass EnumMeta and fiddle with the new Enum's class dictionary before the members are created:
from enum import EnumMeta, Enum, _EnumDict
class StrEnumMeta(EnumMeta):
def __new__(metacls, cls, bases, oldclassdict):
"""
Scan through `oldclassdict` and convert any value that is a plain tuple
into a `str` of the name instead
"""
newclassdict = _EnumDict()
for k, v in oldclassdict.items():
if v == ():
v = k
newclassdict[k] = v
return super().__new__(metacls, cls, bases, newclassdict)
class AutoStrEnum(str, Enum, metaclass=StrEnumMeta):
"base class for name=value str enums"
class Animal(AutoStrEnum):
horse = ()
dog = ()
whale = ()
print(Animal.horse)
print(Animal.horse == 'horse')
print(Animal.horse.name, Animal.horse.value)
Which gives us:
Animal.horse
True
horse horse
1 Disclosure: I am the author of the Python stdlib Enum, the enum34 backport, and the Advanced Enumeration (aenum) library.
Perhaps you are looking for the name attribute which is automatically provided by the Enum class
>>> class Animal(Enum):
... ant = 1
... bee = 2
... cat = 3
... dog = 4
...
>>> Animal.ant.name == "ant"
True
Though if you really want to shoot yourself in the foot. And I'm sure this will introduce a whole world of gotchas (I've eliminated the most obvious one).
from enum import Enum, EnumMeta, _EnumDict
class AutoStrEnumDict(_EnumDict):
def __setitem__(self, key, value):
super().__setitem__(key, key)
class AutoStrEnumMeta(EnumMeta):
#classmethod
def __prepare__(metacls, cls, bases):
return AutoStrEnumDict()
def __init__(self, name, bases, attrs):
super().__init__(name, bases, attrs)
# override Enum.__str__
# can't put these on the class directly otherwise EnumMeta overwrites them
# should also consider resetting __repr__, __format__ and __reduce_ex__
if self.__str__ is not str.__str__:
self.__str__ = str.__str__
class AutoStrNameEnum(str, Enum, metaclass=AutoStrEnumMeta):
pass
class Animal(AutoStrNameEnum):
horse = ()
dog = ()
print(Animal.horse)
assert Animal.horse == "horse"
assert str(Animal.horse) == "horse"
# and not equal to "Animal.horse" (the gotcha mentioned earlier)