Yaroslav Grebnov

How to use enums in Python applications

Posted September 9, 2021 by Yaroslav Grebnov ‐ 16 min read

Introduction

Python enum package allows to utilise very powerful functionality in terms of managing constants and configuration entities. These entities are very compact and in many cases are more convenient to use than dictionaries or configuration files in any format. In a set of examples below, we will demonstrate how to use enums in Python applications as well as the benefits of such usage.

We will start from where we have left off in the How to use constants in Python applications post. As a brief recap, in that post we compared different approaches of using constants as values of an Application class version attibute. Using enums was one of the discussed approaches. In this post, we will develop that basic example by adding more functionality to it.

All examples in this post are connected with each other. They are presented in the complexity increasing order, with examining of the changes made in the code:

  • using a default value in enum,
  • obtaining an enum element from a string,
  • comparing complex enum elements,
  • using chained enums.

In each section, we examine the code of the Application class, the code of the unit tests, and the code of how the class can be used from another module. In unit tests we are using pytest and given-when-then pattern.

The code samples in this post were created and tested using Python 3.8.

Basic example of using enums in Python

In the basic example, we will examine an Application version number implemented as an enum. The version number will follow the Semantic versioning scheme and will consist of three parts: “major.minor.patch”. In the example, we will have three version numbers, “1.0.1”, “1.0.2”, and “2.0.1”. The ApplicationVersion class implementation is the following:

from enum import Enum


class ApplicationVersion(Enum):
    V1_0_1 = "1.0.1"
    V1_0_2 = "1.0.2"
    V2_0_1 = "2.0.1"

    def __init__(self, full):
        self.full = full
        self.major = int(self.full[0])

Notes:

  • the ApplicationVersion class inherits from the enum.Enum class,
  • the version names (like ‘V1_0_1’) are the class attributes. They must follow the rules of Python class attributes naming. That means that we cannot name version “1.0.1” as “1.0.1” because we cannot start an attribute name with a number and cannot use dots in it,
  • we specify version values (like “1.0.1”) in the ApplicationVersion class attributes values,
  • we specify ApplicationVersion accessible properties in the class __init__ method. As we can see from the code, we have two accessible properties, ‘full’ and ‘major’. Both of these two properties values can be retrieved for each version. For example, ‘ApplicationVersion.V1_0_1.full’ or ‘ApplicationVersion.V1_0_2.major’,
  • application version ‘full’ property returns full version number,
  • application version ‘major’ property returns major version number, for example, ‘2’ for ‘2.0.1’.

And now let us proceed to the Application class which uses the ApplicationVersion class for managing application versions data. The Application class has two properties, full_version and major_version returning an application full and major versions accordingly. Application class takes an ApplicationVersion as a parameter on an instance creation. The Application class is implemented as follows:

from enums.application_version import ApplicationVersion


class Application:
    def __init__(self, version: ApplicationVersion):
        self._version = version

    @property
    def full_version(self) -> str:
        return self._version.full

    @property
    def major_version(self) -> int:
        return self._version.major

    def __str__(self):
        return f"Instance of {self.__class__.__name__}, version: {self.full_version}"

Notes:

  • the ‘full_version’ and ‘major_version’ properties are read-only,
  • a __str__ method has been added in order to provide details about the instances in the Application class usage examples.

Let us now write a simple test to check whether the Application properties values are assigned correctly.

import random
from entities.application import Application
from enums.application_version import ApplicationVersion


def test_application_has_expected_version_value():
    # given
    version = random.choice(list(ApplicationVersion))
    # when
    application = Application(version)
    # then
    assert application.full_version == version.full
    assert application.major_version == version.major

Notes:

  • the line version = random.choice(list(ApplicationVersion)) randomly selects an ApplicationVersion object,
  • the randomly selected at the previous step ApplicationVersion object is used to create an Applicatin instance,
  • application full_version and major_version properties are checked to be equal to the ApplicationVersion object corresponding properties values.

As the final part of the basic exmple, we will show how the Application class can be used from other modules. The usage will be very similar to the one already shown in tests.

from entities.application import Application
from enums.application_version import ApplicationVersion


def main():
    application = Application(ApplicationVersion.V1_0_1)
    print(application)


if __name__ == '__main__':
    main()

From the code above, we can see that the Application class usage is straigtforward. Given that the list of the possible ApplicationVersion numbers is visible to the caller, we do not need to have any version values checks before creating an Application class instance.

Using a default value in enum

Let us expand the basic example by adding a new requirement of using a default ApplicationVersion value in case an Application class is initiated without the version argument. Let us consider the 1.0.2 version as a default one.

In order to implement above described requirement, we will add a classmethod with the name default to the ApplicationVersion class. The default method will just return the ApplicationVersion.V1_0_2 object. Such implementation will allow us to return a valid ApplicationVersion object by calling ApplicationVersion.default() without modifying the possible ApplicationVersion numbers list.

@classmethod
def default(cls):
    return cls.V1_0_2

The Application class code changes are minimal. Essentially, we just need to add a default value to the __init__ method version parameter. And naturally, this default value is ApplicationVersion.default(). However, technically there is some semantics involved.

from typing import Optional
from enums.application_version import ApplicationVersion


class Application:
    def __init__(self, version: Optional[ApplicationVersion] = None):
        self._version = version or ApplicationVersion.default()

Notes:

  • version parameter type is changed from ApplicationVersion to Optional[ApplicationVersion],
  • version parameter default value is None,
  • self._version value is equal to either the supplied by caller version argument, if it is not None, or to the ApplicationVersion.default() otherwise,
  • as default is the ApplicationVersion class method, we need to add parenthesis when calling it.

Now we need to add a test checking that if an Application class instance is created without parameters, its full_version and major_version properties are equal to the ApplicationVersion default version object corresponding fulland major properties. Essentially, this test is just a scenario of a test we have already in place. In order to be able to use scenarios in test, we will rearrange the code and utilize the pytest test parametrization functionality.

import random
import pytest
from entities.application import Application
from enums.application_version import ApplicationVersion

random_version = random.choice(list(ApplicationVersion))
test_data = [
    (random_version, random_version),
    (None, ApplicationVersion.default())
]
test_scenarios_ids = [
    "Application initialized with version",
    "Application initialized without version",
]


@pytest.mark.parametrize("input_version,expected_application_version", test_data, ids=test_scenarios_ids)
def test_application_has_expected_version_value(input_version, expected_application_version):
    # when
    application = Application(input_version) if input_version else Application()
    # then
    assert application.full_version == expected_application_version.full
    assert application.major_version == expected_application_version.major

The general idea is very simple, we create an Application class instance with input_version object as parameter. Then, we compare the application properties values with the expected_application_version corresponding properties ones. The implementation, though, needs some explanation:

  • a randomly chosen ApplicationVersion object is assigned to the random_version variable. This variable is evaluated only once,
  • test_data variable holds a list of tuples with pairs of corresponding input_version - expected_application_version objects. For example, the second tuple is equal to (None, ApplicationVersion.default()), which will be processed by the test as Application without version argument value, ApplicationVersion.default(),
  • there are two tuples in test_data which means that there are two test scenarios,
  • test_scenarios_ids variable holds the list of the test scenarios names,
  • @pytest.mark.parametrize decorator glues all the test element together. "input_version,expected_application_version" are the names of the test method arguments. These arguments values are taken from the test_data list. The order of the arguments names in the decorator first argument is the same as the order of values in the test_data list. Test scenarios IDs are taken from the test_scenarios_ids list,
  • the test function has now the input_version and expected_application_version parameters. Earlier, we have discussed them already,
  • in the test function body, all the input version and expected application.version objects are replaced accordingly by the input_version and expected_application_version ones,
  • application = Application(input_version) if input_version else Application() line calls Application with version=input_version if the latter exists (or not None in our case) or just Application() otherwise.

As the final part of the example, let us see how the Application class instance can be created without version argument from another module.

def main():
    application = Application()
    print(application)

As expected, the code above prints out the details of an Application class instance with full_version property returning the default ApplicationVersion object.

Instance of Application, version: 1.0.2

Obtaining an enum element from a string

Let us examine a frequent use case of initializing classes using values retrieved from environmental variables. In this example, the Application version value will be taken from the APPLICATION_VERSION environment variable. As an environment variable can hold a value only of type string, we need to convert that string value into an enum element before passing to the Application class constructor. The conversion will be implemented as the ApplicationVersion classmethod. Please note that in this case the input is a string argument which can have any value. In order to make sure that this supplied value is correct, we need to have a corresponding check in the method.

@classmethod
def from_string(cls, version_name: str):
    try:
        return cls[f"V{version_name.replace('.', '_')}"]
    except KeyError:
        raise ValueError(f"'{version_name}' is not a valid application version name. "
                         f"Please use one of: {[e.value for e in cls]}")

Notes:

  • the from_string method implementation is based on the fact that an Enum class has a very convenient additional form of presentation as a dictionary with the enum elements names as keys and enum elements objects as values. This gives us a possibility to get, for example, an ApplicationVersion.V1_0_1 object by calling ApplicationVersion['V1_0_1'],
  • we will be getting ApplicationVersion objects by their names in “n.n.n” format, for example, “1.0.1”. In order to get the real names, we need to perform a basic transformation: add a ‘V’ prefix and replace dots by underscores,
  • the method body is wrapped into a try-catch block. The purpose of this is to catch errors in cases of incorrect version_name string values provided as argument when calling the method. Please note that internally the possible error is of type KeyError risen in case of an incorrect dictionary key usage. But for the external caller, the error is of type ValueError meaning that the ApplicationVersion element value is incorrect,
  • ValueError message contains a list of possible application version numbers: f"Please use one of: {[e.value for e in cls]}".

The Application class code is not modified in this example.

In order to test the newly added from_string method, we will create two test functions. The first one will be used to check an ApplicationVersion object creation from an existing application version number supplied as a string. The second one will check if an error is risen in case a non-existing application version number is supplied as a string.

def test_application_version_is_correctly_initialized_from_string():
    # given
    random_application_version = random.choice(list(ApplicationVersion))
    random_application_version_number = random_application_version.value
    # then
    assert ApplicationVersion.from_string(random_application_version_number) == random_application_version


def test_application_version_raises_error_on_incorrect_string_version_number():
    # given
    version_number = "1.0.5"
    # then
    with pytest.raises(ValueError):
        ApplicationVersion.from_string(version_number)

Notes on the first test code:

  • in the given part first line, we randomly choose an ApplicationVersion object from a list of possible ones,
  • in the given part second line, we note the randomly chosen ApplicationVersion object string value,
  • the subsequent checck is a simple assert.

Notes on the second test code:

  • as we do not know how to compare the ApplicationVersion objects yet, we cannot programmatically retrieve the boundary objects from the list of possible objects. This leaves us no other choice except of hardcoding a non-existing application version number,
  • with pytest.raises(ValueError): line allows us to check if a ValueError is risen on calling ApplicationVersion.from_string method with non-existing application version number.

Finally, let us analyse an example of how adding the functionality of creating ApplicationVersion objects from strings impacts the usage of Application class from other modules. As it has been mentioned at the beginning of the section, application version numbers are retrieved from environment variables. These numbers are converted from strings into ApplicationVersion objects by using its from_string method. As the string numbers can potentially be non-existing application version numbers which can lead to throwing a ValueError, we wrap the conversion and Application instance creation in a try-catch block.

import os
from entities.application import Application
from enums.application_version import ApplicationVersion


def main():
    try:
        for version_number in {os.getenv("CORRECT_APPLICATION_VERSION"), os.getenv('INCORRECT_APPLICATION_VERSION')}:
            version = ApplicationVersion.from_string(version_number)
            application = Application(version)
            print(f"Initialized an Application: {application}")
    except ValueError as e:
        print(f"Error initializing an Application: {e}")


if __name__ == '__main__':
    os.environ["CORRECT_APPLICATION_VERSION"] = "1.0.1"
    os.environ["INCORRECT_APPLICATION_VERSION"] = "1.0.5"
    main()

Notes:

  • we are using two environment variables, “CORRECT_APPLICATION_VERSION” and “INCORRECT_APPLICATION_VERSION”, containing correct and incorrect application version numbers accordingly,
  • two Application class initialization examples are implemented as a for-loop, looping over a set of correct and incorrect application version environment variables values,
  • in case the ApplicationVersion.from_string method throws an error, it is wrapped by an Application class initialization error message.

The code above outputs the following:

Initialized an Application: Instance of Application, version: 1.0.1
Error initializing an Application: '1.0.5' is not a valid application version name. Please use one of: ['1.0.1', '1.0.2', '2.0.1']

Comparing complex enum elements

Let us consider further extension of the ApplicationVersion class functionality by adding a new requirement. We need to be able to compare application versions by applying the following rule: version “x1.y1.z1” is greater than version “x2.y2.z2” if and only if x1 > x2 or (x1 == x2 and y1 > y2) or (x1 == x2 and y1 == y2 and z1 > z2). By design, we cannot have two different ApplicationVersion objects with the same application version number. That means that in order to compare ApplicationVersion objects, we need to implement only one rule stated above.

In order to implement the new requirement, we need to add four new methods to the ApplicationVersion class: __lt__, __le__, __gt__, and __ge__. As we have mentioned above, the ApplicationVersion class design allows us to simplify this task and implement only one method, either __lt__ or __gt__. All other methods implementations will derive from the first one.

def __gt__(self, other):
    return [int(x) for x in self.full.split(".")] > [int(x) for x in other.full.split(".")]

def __ge__(self, other):
    return self == other or self > other

def __lt__(self, other):
    return other > self

def __le__(self, other):
    return other >= self

Notes:

  • we have chosen to implement the __gt__ method. Its implementation is based on the Python’s collections comparison functionality, meaning that each element of the fist collection is compared with the corresponding (by order) element of the second collection, which conveniently satisfys the condition described in the rule at the beginning of the section,
  • [int(x) for x in self.full.split(".")] (and the same for other ApplicationVersion) creates a list of converted into integers “x”, “y”, “z” elements of version number “x.y.z”,
  • all subsequent comparison methods are derived from the methods derived before them.

There is one more note which we would like to emphasize. Implementation of the comparison methods gives us an additional benefit of possibility to get maximum and minimum application versions without adding any additional line of code. We will demonstrate this later in tests.

The Application class code is not modified in this example.

In order to test our implementation of the new requirement, two new test functions are added: the first - for testing ApplicationVersion comparisons, the second - for testing correct determination of minimum and maximum ApplicationVersion. Both tests are simple and only contain lists of assertions. In order to demonstrate different comparison cases, we have added some additional ApplicationVersion objects:

    V1_0_1 = "1.0.1"
    V1_0_2 = "1.0.2"
    V1_30_7 = "1.30.7"
    V2_0_1 = "2.0.1"
    V2_5_1 = "2.5.1"
    V12_0_1 = "12.0.1"

The new tests code:

def test_application_versions_are_correctly_compared():
    assert ApplicationVersion.V1_0_1 == ApplicationVersion.V1_0_1
    assert ApplicationVersion.V1_0_1 != ApplicationVersion.V1_0_2
    assert ApplicationVersion.V1_0_2 > ApplicationVersion.V1_0_1
    assert ApplicationVersion.V2_0_1 >= ApplicationVersion.V1_0_2
    assert ApplicationVersion.V1_30_7 < ApplicationVersion.V2_5_1
    assert ApplicationVersion.V2_0_1 <= ApplicationVersion.V12_0_1
    assert not ApplicationVersion.V1_0_1 > ApplicationVersion.V1_0_1
    assert not ApplicationVersion.V1_0_1 < ApplicationVersion.V1_0_1


def test_min__max_application_versions_are_correctly_identified():
    assert max(ApplicationVersion) == ApplicationVersion.V12_0_1
    assert min(ApplicationVersion) == ApplicationVersion.V1_0_1

Given that now we can compare the ApplicationVersion objects, we can update one of the tests from the previous section. The hardcoded incorrect application version number value can be replaced by the maximum ApplicationVersion incremented by one. The updated test code will look like:

def test_application_version_raises_error_on_incorrect_string_version_number():
    # given
    version_number = f"{max(ApplicationVersion).major + 1}{max(ApplicationVersion).full[1:]}"
    # then
    with pytest.raises(ValueError):
        ApplicationVersion.from_string(version_number)

The code of the Application class usage from other modules is not modified in this example.

Using chained enums

In the final example of this article, we will show how to use chained enums. This use case is quite frequent in practice. However, we would like to note that implementation of such use cases requires great attention to details while designing the enums as it may be relatively easy to create cyclic dependencies.

In this example, we will have a new requirement. Depending on its version number, an application uses a different type of database: an application with major_version equal to “1” uses SQLite, the one with higher major_version - PostgreSQL. While installing an application, we need to execute a specific script checking the database setup. This script is different for each database type.

First of all, we will create a new DatabaseType enum to hold the corresponding data. As per requirements above, we have an SQLITE and a POSTGRESQL database types. script_path property returns path to the corresponding database type setup script.

from enum import Enum


class DatabaseType(Enum):
    SQLITE = "/path/to/sqlite-script"
    POSTGRESQL = "/path/to/postgresql-script"

    def __init__(self, script_path):
        self.script_path = script_path

In the ApplicationVersion class, we will implement the mapping between an application version and a database type. We will add a database_type property to ApplicationVersion. As per requirements, this property will return DatabaseType.SQLITE for ApplicationVersions with major version number equal “1” and DatabaseType.POSTGRESQL otherwise.

from enum import Enum
from enums.database_type import DatabaseType


class ApplicationVersion(Enum):
    V1_0_1 = "1.0.1"
    V1_0_2 = "1.0.2"
    V1_30_7 = "1.30.7"
    V2_0_1 = "2.0.1"
    V2_5_1 = "2.5.1"
    V12_0_1 = "12.0.1"

def __init__(self, full):
    self.full = full
    self.major = int(self.full[0])
    self.database_type = DatabaseType.SQLITE if self.major == 1 else DatabaseType.POSTGRESQL

Notes:

  • we are using a one-line form of the if-else statement.

From the Application class point of view, we are interested only in the database script setup path. In order to be able to retrieve it, we will add a corresponding property to the Application class. The database_script_path property will return the path to the database setup script after retrieving it from the _version internal variable value through a chain of enum objects.

@property
def database_script_path(self) -> str:
    return self._version.database_type.script_path

In order to test correctness of the Application class database_script_path property value, we will add a corresponding test function. The new test will be parametrized and will have two scenarios: one for an application with “1.y.z” version which should have database_script_path pointing to DatabaseType.SQLITE script path, and another one for an application with “not 1.y.z” version which should have database_script_path pointing to DatabaseType.POSTGRESQL script path.

Implementation of the test is similar to the one for the versions from strings. As we have discussed the latter in details already, we will just provide the new test code.

test_data = [
    (ApplicationVersion.V1_0_1, DatabaseType.SQLITE.script_path),
    (ApplicationVersion.V2_0_1, DatabaseType.POSTGRESQL.script_path)
]
test_scenarios_ids = [
    "SQLite script_path for 1.y.z application version",
    "PostgreSQL script_path for >1.y.z application version",
]


@pytest.mark.parametrize("input_version,expected_database_script_path", test_data, ids=test_scenarios_ids)
def test_application_has_expected_script_path_value(input_version, expected_database_script_path):
    # when
    application = Application(input_version)
    # then
    assert application.database_script_path == expected_database_script_path

We will update the Application class from other module usage example with getting an Application instance database_script_path property value. The changes are minimal and self-explanatory.

def main():
    try:
        version = ApplicationVersion.from_string(os.getenv("APPLICATION_VERSION"))
        application = Application(version)
        print(f"Initialized an application: {application}")
        print(f"Executing the application database setup check script: {application.database_script_path}")
    except ValueError as e:
        print(f"Error initializing an application: {e}")


if __name__ == '__main__':
    os.environ["APPLICATION_VERSION"] = "1.0.1"
    main()

The code outputs:

Initialized an application: Instance of Application, version: 1.0.1
Executing the application database setup check script: /path/to/sqlite-script

Conclusion

We have discussed in detail five examples of how enums can be used in a Python application. We have started from the most simple example and proceeded with adding new functionality in the subsequent ones. We have discussed the cases of a default enum object, obtaining an enum object from a string, comparing complex enum objects, and chaining the enums. The examples use a class with properties based on enums. In each section, we provide tests samples as well as how to use the above mentioned class from another module. Each example is discussed in detail and has line-to-line code explanations.

In case you have a question or a comment concerning this post, please send them to: