An alternative to Enum for Choices

Published:

Those who've been reading my older posts may remember I showed how you could use Enum and IntEnum as a cleaner way to declare const-type values for choices lists in Django fields.

That solution never felt comfortable to me, because Enum values aren't simple values.

So after some playing around, and a brief look over the enum.py in python 3.5, I've come up with the following:

class ChoiceProperty:
    '''Descriptor class for yielding values, but not allowing setting.'''
    def __init__(self, value):
        self.value = value

    def __get__(self, instance, cls=None):
        return self.value


class MetaChoices(type):
    @classmethod
    def __prepare__(mcs, name, bases, **kwargs):
        '''Use an ordered dict for declared values.'''
        return OrderedDict()

    def __new__(mcs, name, bases, attrs):
        _choices = OrderedDict()
        _label_map = {}

        for name, value in list(attrs.items()):
            if not name.isupper():
                continue
            if isinstance(value, tuple):
                value, label = value
            else:
                label = name.title().replace('_', ' ')
            _choices[value] = label
            _label_map[label] = value
            attrs[name] = ChoiceProperty(value)
        attrs['_choices'] = _choices
        attrs['_labem_map'] = _label_map

        return type.__new__(mcs, name, bases, dict(attrs))

    def __getitem__(cls, key):
        return cls._choices[key]

    def __iter__(cls):
        return iter(cls._choices.items())


class Choices(metaclass=MetaChoices):
    '''Base class for choices constants.'''

So, what does this all mean?

First, there's the ChoiceProperty. This follows the descriptor property so we can have attributes on our class that simply return a value they were told. They do this irrespective of if it's on an instance or the class itself!

Last is the Choices class, which is empty except for declaring it has a metaclass.

In the middle is the meat of the work, of course. A Metaclass defines what is done when you declare a class, or a subclass. This lets you, as you can see here, iterate everything you're declaring on the class and do something with it before the class is declared.

So in this case it's finding all attributes of the class whose name is SHOUTY_SNAKE_CASE, and treating them as const declarations.

Either they're NAME = Value, and a label is created from the NAME, or they're NAME = Value, Label.

The __getitem__ method is called when you try to subscribe the class (i.e. FOO[0]).

And the __iter__ method when you try to iterate it.

So what can I do with it?

>>> class STATE_CHOICES(Choices):
...     NEW = 0
...     IN_PROGRESS = 1
...     REVIEW = 2, 'In Review'
...
>>>
>>> STATE_CHOICES.NEW
0
>>> STATE_CHOICES.IN_PROGRESS
1
>>> STATE_CHOICES[2]
'In Review'
>>> list(STATE_CHOICES)
[(0, 'New'), (1, 'In Progress'), (2, 'In Review')]

Now, does that look useful?

class MyModel(models.Model):
    class STATUS(Choices):
        CLOSED = 0
        NEW = 1
        PENDING = 2, 'Process Pending'
        FAILED = -1, 'Processing Failed'

    status = models.IntegerField(choices=list(STATUS), default=STATUS.NEW)