A C style enum in Python
The enum standard module allows you to make names whose values are arbitrary:
from enum import Enum
class Color(Enum):
RED = 0
GREEN = 1
BLUE = 2
It is convenient to use Color.RED in code because the compiler/interpreter will ensure against typos, and sensibly chosen names will enhance readibility. However, unlike in C, Python’s enum forces you to assign values. It would be nice if we can simply write
class Color(Enum):
RED
GREEN
BLUE
and have the attributes autoincrement. It turns out that with a little metaprogramming, we can.
Python 3 introduced __prepare__, which is a method on the metaclass invoked before class creation. It by default returns a dictionary, which will contain the class attributes required for class creation. For example, the statement RED = 1 above has the effect of adding a key/value pair to this dictionary: prepared_dict['RED'] = 1.
The idea is to override __prepare__ and return a customized dictonary that assigns each key the next available integer. Here is an implementation of this dictionary:
class AutoDict(dict):
def __init__(self):
self.count = 0
def __getitem__(self, key):
if key.startswith("__") and key.endswith("__"):
return super().__getitem__(key)
self[key] = self.count
self.count += 1
Under a class definition, for each variable that appears in an expression - as opposed to the left hand side of an assignment - Python will look up its value first in this dictionary, calling the __getitem__ method, which in turn assigns to it the next integer.
Then we create a metaclass with a __prepare__ method that returns an instance of AutoDict,
class EnumMeta(type):
def __prepare__(*args):
return AutoDict()
and a base class that instantiates EnumMeta
class Enum(metaclass=EnumMeta):
pass
Let’s try running the example above:
>>> class Color(Enum):
... RED
... GREEN
... BLUE
>>> Color.RED
0
>>> Color.GREEN
1
>>> Color.BLUE
2
So we have an enum class that automatically assign and increment values to its attributes.