Last Updated: February 25, 2016
·
3.397K
· ruiwen

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