Fun With Python Enums

Enums have been around in Python since version 3.4, released in 2014, but I still don’t feel like they get enough love.

An Enum is a type that stores a number of specific related constant values. Any value not defined on the Enum is invalid and unable to be stored.

For newbies, this will post will get you up to speed on Enums, but I’ve also included a couple of neat Enum tricks so even experienced developers should be able to learn something.

What is an Enum and what do they do?

Enums inherit from the enum.Enum class, and the attributes are defined like standard Python class attributes. They can have any value you choose – as long as each one is unique.

For simplicity you can use the enum.auto() function to automatically define these attributes. When doing this, their values will be successive integers starting at 1.

The classic example is an Enum that stores compass directions, North, South, East and West.

from enum import Enum, auto

class Cardinal(Enum):
    NORTH = auto()
    SOUTH = auto()
    EAST = auto()
    WEST = auto()

In this case, the value of Cardinal.NORTH is 1, Cardinal.SOUTH is 2, and so on – although the point of Enums is that we don’t need to care about what these values actually are.

Enum variables are set from the Enum’s class attributes.

direction = Cardinal.NORTH

if direction == Cardinal.NORTH:
    print("Facing north.")

This of course helps with autocompletion in your IDE and prevents being able to use incorrect (missing) values. Consider if we tried to store the cardinal as a string instead.

direction = "north"

if direction == "NORTH":
    print("Whoops! We forget this should be lowercase!")

Every time you wanted to do a comparison like this, you would have to remember if it was North, NORTH or north, and your IDE would be no help here.

You could attempt to emulate the behaviour of Enums with a standard class, but you lose some of the helper methods that Enums provide. For example, you can get the name of a particular Enum attribute:

>>> west = Cardinal.WEST
>>> west.name
'WEST'

You can also get its value. When using Enums with automatic attributes this isn’t very useful:

>>> west.name
4

But, this comes in more handy when defining your own values on an Enum. You can reference the value, and use the Enum to validate/clean “dirty” values.

In this next example, the Enum represents a URL scheme, which can be either http or https.

class Scheme(Enum):
    HTTP = "http"
    HTTPS = "https"

Once we have a Scheme reference, we could use it to build a URL:

>>> scheme = Scheme.HTTP
>>> host = "example.com"
>>> path = "/document.html"
>>> url = f"{scheme.value}://{host}{path}"
>>> url
'http://example.com/document.html'

We can also create a Scheme instance from a string, by passing the enum value to the constructor.

>>> Scheme("http")
<Scheme.HTTP: 'http'>

This cleans the value, by transforming the string into the Scheme type. It also provides validation, as a ValueError will be raised if trying instantiate a Scheme (or other Enum) from a value that it doesn’t have.

>>> ftp = Scheme("ftp")
ValueError: 'ftp' is not a valid Scheme

Enums also provide protection for trying to assign the same value to multiple attributes – this is invalid:

class BadEnum(Enum):
    VALUE = 1
    OTHER_VALUE = 1  # bad, duplicate value

That’s a brief intro but covers 99% of the day to day use cases of Enums.

Let’s now have a look at when you should use them.

Better Booleans

Next time you’re writing a function that accepts a boolean, consider if it really is a boolean or are you trying to represent one of two opposite values? And, could there be other mutually exclusive values in this set?

To elaborate with an example, let’s say we want to retrieve a list of values, sorted either ascending or descending.

The function might look like this:

def fetch_sorted_values(sort_ascending: bool) -> List[int]:
    values = load_values()  # read from file etc
    values.sort(reverse=not sort_ascending)
    return values

This function works OK, but could be better. Consider the code when we call the function.

values = fetch_sorted_values(True)

Pop quiz: without referring to the function signature, in what order are the values sorted?

Also, let’s say that later we want to enhance the function to also return values in a random order. If we add another argument (sort_random), then what happens if both sort_ascending and sort_random are True? The function would be even more confusing to call, too:

values = fetch_sorted_values(False, True)  # what does this mean?!

Enter Enums to save the day. First define a simple one to indicate the sort direction:

class SortDirection(Enum):
    ASCENDING = auto()
    DESCENDING = auto()

Then implement it inside fetch_sorted_values.

def fetch_sorted_values(sort_direction: SortDirection) -> List[int]:
    values = load_values()  # read from file etc
    values.sort(reverse=sort_directory == SortDirection.DESCENDING)
    return values

Now when the function is called it’s a lot more obvious what the behaviour is.

values = fetch_sorted_values(SortDirection.ASCENDING)

We can also alter the Enum to add the random sorting option:

class SortDirection(Enum):
    ASCENDING = auto()
    DESCENDING = auto()
    RANDOM = auto()

Then just alter the function implementation without changing the signature – the implementation is left as an exercise to the reader.

Command Line Arguments

Enums can easily be used with argparse, for parsing command line arguments to your script.

The ArgumentParser.add_argument method takes two parameters for validating/cleaning command line arguments, type and choices.

The type argument is a function/class that will be called to convert the string to a Python type. This could be int to convert from str to int, for example.

choices takes an iterable (list, tuple, etc) of values that are valid for the argument. It is checked after the arguments have been converted with the type argument function, so should contain values of that type. For example, for int types, pass in the choices [1, 2, 3] not ["1", "2", "3"].

We can pass our Enum class into these arguments. Not only will this validate that the given value is valid for the Enum, but it will also convert the argument to the Enum type.

The only caveat for this is that the Enum values must be strings, for the automatic conversion to work, as the arguments coming in from the command line (sys.argv) will be strings too.

Here’s a simple example, a Python script that controls another external process. It will take just one argument on the command line, the action to take on the process. This will either be start to start the process or stop to stop it.

class ControlAction(Enum):
    START = "start"
    STOP = "stop"

Then with argparse, we can automatically parse the value.

We just set the argument’s type to the ControlAction class, which will make the argument be converted from a str to ControlAction.

We also pass ControlAction as the choices parameter. This is sort of redundant, since the conversion from str to ControlAction will fail anyway if an invalid value is supplied, however it does mean that a list of valid values will be automatically printed out if one is not supplied.

parser = argparse.ArgumentParser(
    description="Control an external process"
)

parser.add_argument(
    "control_action",
    action="store",
    type=ControlAction,
    choices=ControlAction
)

args = parser.parse_args()

print(args.control_action)  # this will be a ControlAction

Now here is how the script responds to valid and invalid arguments:

$ python3 process_controller.py start
ControlAction.START

$ python3 process_controller.py stop
ControlAction.STOP

$ python3 process_controller.py bad
usage: process_controller.py [-h] {
    ControlAction.START,ControlAction.STOP
}
process_controller.py: error:
    argument control_action: invalid ControlAction value: 'bad'

You can see how the string value is automatically converted to an Enum and validated.

Also remember that if we wanted to add an extra action it’s trivial. It just needs to be added to the ControlAction Enum and then the value is automatically parsed and validated.

Of course you still need to write code to undertake the new, action but that’s a given in any case.

Other Cool Things

Enums are iterable:

>>> list(Scheme)
[<Scheme.HTTP: 'http'>, <Scheme.HTTPS: 'https'>]

Python 3.6 onwards also includes a type of enum called IntFlag. While I won’t go into depth here since it would basically be duplicating the official documentation, it’s worth pointing out.

It lets you define Enum values as integers, then perform bitwise operations on them. If you’re familiar with Unix permissions then this code snippet should make sense:

>>> class Perm(IntFlag):
...     R = 4
...     W = 2
...     X = 1
...     RWX = 7
...
>>> Perm.R | Perm.W | Perm.X
<Perm.RWX: 7>

It combines the permission values and then converts that to its own Enum value, if one exists. Also check the Flag class for a similar yet different type that may more closely match your needs for bitwise Enum comparison.

Conclusion

As I wrote at the start, despite being around for a while I don’t see Enums being used very much and I think they deserve more love. Their ability to automatically validate and convert data can cut down on boilerplate code, as well as make your code easier to write as your IDE will be more likely to understand them than just plain strings or ints.

While this post was just a brief introduction, it should cover 99% of day to day Enum use cases. Using Enums can also make your code more readable and more flexible to changing requirements later on.