Tuesday 14 May 2013

Python's new Enum class

I used to wish to have an enumeration in Python, like you have in other languages.

On May 10th, GvR accepted PEP 435, which proposes the addition of an Enum type to the standard library.

Before this, us python coders would need to use plain classes for enums, which doesn't work so well. We can't get enum names by value, for example.

It follows that it's not possible to create a list of 2-tuples for use with Django models. We'd need to have a class mapping our choice names to their values, but you can't create the good old "choices" structure, so that's pretty useless. However, since our new enums are iterable we can:

    class ChoiceClass(IntEnum):
        foo = 1
        bar = 2
        baz = 3

    CHOICES = [(e.value, e.name) for e in ChoiceClass]

Creating your enumerations

By inheriting from Enum, you can create your own enumerations. If all your enumeration values are supposed to be of the same type, you should inherit Enum as well as that type.

    class Numbers(int, Enum):  # you can also use IntEnum
        one = 1
        two = 2
        three = 3

    class Text(str, Enum):
        one = 'one'
        two = 'two'
        three = 'three'

Using your enumerations

By iterating over your enum class, you get the enumeration members, which can be queried for their name or value.

    print(' '.join(['{}: {}'.format(number.name, number.value) for number in Numbers]))

Internally, there is some metaclass magic going on, to allow us to succintly declare these enums. The class comes with a few facilities, like __iter__ as I showed above, __call__ which gets enumeration items by value, and __getitem__ which gets them by name.

You should understand that enumeration items are actually instances of your enumeration class. Which allows us to say isinstance(Text.one, Text).

You can get your enumeration items in several ways.

  • As an attribute: Numbers.one, the Java way
  • By name: Numbers['one'], a tad prettier than using getattr
  • By value: Numbers(1), as a constructor

In a way, the third syntax reminds me of how we convert values using their constructors in python, for instance, using str or int to convert existing objects to string or integer values.

You can start using these enums right away. There was talk of porting this back to python 2 on the python-dev mailing list, but right now the code is for python 3 only. I'm going to convert some Django code to use these enums later, because I have three declarations for each Choice I need (SOMETHING_CHOICES, SOMETHING_DICT, SOMETHING_REVERSE), which is just backwards.

You can grab a copy of this module at bitbucket. Beware, because this was only put up for PEP evaluation. You'll want to wear a helmet and keep in mind that this is not stable or supposed to be used in production code. Although Guido has accepted it at the time of writing, it may still be subject to change, since the PEP is not in the "final" status.

That said, grab the "ref435" module from that package and try it out.

    >>> from ref435 import Enum
    >>> class Enumeration(int, Enum):
    ...     potatoes = 1
    ...     apples = oranges = 2
    ... 
    >>> Enumeration.apples is Enumeration.oranges
    True

You can look for reference in the PEP.

See you next time!

9 comments: