Enums are cool after all

It took me some time to realize that Enums were a cool thing in Python. Before that, I thought that I could achieve the same result with a bunch of constants and dicts. In my current project, I found myself using them all over the place.

Enum exists in a lot of languages, from C to Java. It is basically a set of possible integer values, labeled with a name. For instance, Spring, Summer, Autumn and Winter could be part of an enum called Season.

A specific structure is not necessary to do so, as you can achieve the same goal with simple integer constants. But the enum cares about the boring stuff for you. Quick tour.

Import the Enum base class to start playing with Enum (I leave the ipython beginning of lines)

In [1]: from enum import Enum

You can then create your first Enum

In [2]: class Season(Enum):
...:     SPRING = 0
...:     SUMMER = 1
...:     AUTUMN = 2
...:     WINTER = 3

And instantiate a member

In [3]: s = Season.SUMMER

You can access a string representation of the enum member, its value and name

In [4]: s
Out[4]: <Season.SUMMER: 1>

In [5]: str(s)
Out[5]: 'Season.SUMMER'

In [6]: s.name
Out[6]: 'SUMMER'

In [7]: s.value
Out[7]: 1

Enums are good to compare values in a readable way

In [8]: s == Season(1)
Out[8]: True

Another cool stuff is the ability to return all enum possible value as a list

In [9]: list(Season)
Out[9]:
[<Season.SPRING: 0>,
<Season.SUMMER: 1>,
<Season.AUTUMN: 2>,
<Season.WINTER: 3>]

You are not obliged to set the integer values manually as enum package provides auto function to do it for you, starting with 1.

In [10]: from enum import auto

In [11]: class Season2(Enum):
...:     SPRING = auto()
...:     SUMMER = auto()
...:     AUTUMN = auto()
...:     WINTER = auto()
...:

In [12]: Season2.SPRING
Out[12]: <Season2.SPRING: 1>

In [27]: list(Season2)
Out[27]:
[<Season2.SPRING: 1>,
<Season2.SUMMER: 2>,
<Season2.AUTUMN: 3>,
<Season2.WINTER: 4>]

You can start with the value you want, then use auto to increment it.

In [22]: class Season4(IntEnum):
...:     SPRING = 0
...:     SUMMER = auto()
...:     AUTUMN = auto()
...:     WINTER = auto()


In [24]: list(Season4)
Out[24]:
[<Season4.SPRING: 0>,
<Season4.SUMMER: 1>,
<Season4.AUTUMN: 2>,
<Season4.WINTER: 3>]

Unfortunately, you cannot compare an enum member with an int directly

In [13]: Season2.SPRING == 1
Out[13]: False

In [25]: Season2.SUMMER == 2
Out[25]: False

But you can use IntEnum alternative to do so. Its an enum that subclasses the int class to ease comparison and affectation. It can be a good idea if you use enum along with integer fields in a database.

In [14]: from enum import IntEnum

In [15]: class Season3(IntEnum):
...:     SPRING = auto()
...:     SUMMER = auto()
...:     AUTUMN = auto()
...:     WINTER = auto()

In [26]: Season3.SUMMER == 2
Out[26]: True

Finally, as Enum is a class, you can add methods to customize instantiation and representation.

I like using enums to define integer choices for db fields in Django, thanks to the list representation of enum members. It helps me to keep the code base readable and add namespacing rather than using individual constants.

Posted on 2019-09-17 at 23:00

Tags: programming

Previous Back Next