Easy enums in Python
I was recently looking for a way to create simple Enums in Python, and I realised while enums are part of the standard library in Python 3.4, and there have been backports and various other implementations for Python 2.x, they didn't really suit my needs.
I was looking for something
a) that mapped easily to integer values, eg.
>>> user.type == User.TYPES.NORMAL == 1
True
b) that allowed me to customise the sequence progression for items in the enum, eg.
>>> UserEnum = enum("NORMAL", "BUYER", "SELLER", "ADMIN")
>>> UserEnum.NORMAL
0
>>> UserEnum.BUYER
1
>>> UserEnum.SELLER
2
>>> UserEnum.ADMIN
4
Having a powers-of-2 (ie. [0, 1, 2, 4]. I know 0 technically isn't a power of 2, but we can use it to indicate the base state) progression is really handy since we can treat the enum field essentially as a bitfield, ie.
>>> user.type
3
>>> (user.type & UserEnum.SELLER) == UserEnum.SELLER
True
>>> (user.type & UserEnum.BUYER) == UserEnum.BUYER
True
With these two requirements in mind, the following snippet was developed with help from a couple of StackOverflow posts on the topic
def enum(*sequential, **named):
'''
Easy enums in Python 2.x
Borrowed from http://stackoverflow.com/a/1695250
Metaclass idea from http://stackoverflow.com/a/6970831
Enum values map to integers in a powers-of-2 scale, ie. 0, 1, 2, 4, 8 etc.
Enum class is also hardened against modification through overriding of the
__setattr__ and __setitem__ methods
>>> Numbers = enum('ZERO', 'ONE', 'TWO')
>>> Numbers.ZERO
0
>>> Numbers.ONE
1
'''
enums = dict(zip(sequential, map(lambda x: 2 ** (x - 1) if x > 0 else 0, range(len(sequential)))), **named)
reverse = dict((value, key) for key, value in enums.iteritems())
def not_implemented():
raise NotImplementedError()
meta = type('EnumMeta', (type,), {
"__contains__": lambda self, v: v in enums.keys() or v in enums.values(), # self.__dict__,
"__len__": lambda self: len(sequential),
"__getitem__": lambda self, k: self.__dict__[k] if k in self.__dict__ else reverse[k],
"__setattr__": lambda self, k, v: not_implemented(),
"__setitem__": lambda self, k, v: not_implemented(),
"__iter__": lambda self: iter(sequential)
})
d = {}
d.update(enums)
return meta('Enum', (), d)
In order to avoid accidental/malicious overriding of enum values, we also nerf the magic methods __setattr__
and __setitem__
, giving
>>> UserEnum.BUYER
1
>>> UserEnum.BUYER = 4
Traceback (most recent call last):
...
NotImplementedError
>>> UserEnum['BUYER'] = 4
Traceback (most recent call last):
...
NotImplementedError
>>> setattr(UserEnum, 'BUYER', 4)
Traceback (most recent call last):
...
NotImplementedError
Hope that's been useful. If there's a way that this snippet can be improved, please let me know!
Here's a gist