How To Implement Enumerations With The Python Enum Module

Posted on by Robert Charles
Python Enumerations

C/C++ programmers must be familiar with the enum data type. It allows them to map certain names to unique integers. This tutorial covers the Python enum module, which is designed to provide similar capabilities.

What Is The Use Of Enumerations?

An enumeration is a complete collection of symbolic names that are assigned to unique values (typically integers). It can be seen as a user-defined data type whose names and values are given by the programmer.

Enumerations can come in handy when you work with variables that can only accept a small number of related values. Common examples are days in a week or seasons in a year. Enumerations can make your code more readable and clearer when they are properly defined with meaningful names.

Is There Enumeration In Python?

Traditionally, Python didn’t support enumerations out of the box like C/C++. It wasn’t until the 3.4 version that the enum module was added. It comes with the Enum class, providing a solution for creating enumerations in Python.

What Are Enumerations In Python?

Instances of the Enum class are sets of symbolic names bound to unique values. You can create a class based on it to define a new enumeration, its members, and associated values.

Create An Enumeration With The Python enum Module

You can use the class syntax to create an enumeration:

from enum import Enum
class SocialNetwork(Enum):
        TWITTER = 1
        TIKTOK = 2
        FACEBOOK = 3
        INSTAGRAM = 4

In this example, we have imported the Enum class from the enum module. The SocialNetwork enumeration is created with a class definition, inheriting the Enum class.

In the body of this definition, we define all the members and their values. You should use upper-case letters to name these members, whose names must be unique.

The function-call syntax can also be used to create the above enumeration by deriving from the Enum class:

SocialNetwork = Enum('SocialNetwork', ['TWITTER', 'TIKTOK', 'FACEBOOK', 'INSTAGRAM'])

By default, Python will assign integer values to these names, starting from 1. The values of members aren’t always important – it depends on the nature of your enumeration.

By default, the enum module allows you to bound multiple names to the same value. For instance, this is a totally valid enumeration:

class SocialNetwork(Enum):
        TWITTER = 1
        TIKTOK = 2
        FACEBOOK = 3
        INSTAGRAM = 1

Since the member TWITTER is defined first, INSTAGRAM will be treated as an alias of it. If you look for members by value 1, Python will return the first member (TWITTER). This is also the case even when you do a by-name lookup for INSTAGRAM.

>>> SocialNetwork(1)
<SocialNetwork.TWITTER: 1>
>>> SocialNetwork.INSTAGRAM
<SocialNetwork.TWITTER: 1>

To make sure that every member in an enumeration has a unique name, use the unique() decorator. It will raise errors when there are duplicate member names:

from enum import Enum, unique
@unique
class SocialNetwork(Enum):
        TWITTER = 1
        TIKTOK = 2
        FACEBOOK = 3
        INSTAGRAM = 1

Output

ValueError: duplicate values found in <enum 'SocialNetwork'>: INSTAGRAM -> TWITTER

Remember that while reusing values is possible, this isn’t the case with member names. They always have to be unique, even without the unique() decorator:

class SocialNetwork(Enum):
        TWITTER = 1
        TIKTOK = 2
        TWITTER = 3
        INSTAGRAM = 4

Output

TypeError: Attempted to reuse key: 'TWITTER'

When you don’t care about the exact values of your enumeration’s members, you can use the auto() method to generate them for you:

from enum import Enum, auto
class Language(Enum):
        PYTHON = auto()
        JAVA = auto()
        C = auto()
        SCALA = auto()

Access Members Of Enumerations

You can use its name or value to access a member in an enumeration. The name can be used with the dot operator or indexing operator (square brackets) and values with parentheses (the call syntax).

>>> print(SocialNetwork.TWITTER)
SocialNetwork.TWITTER
>>> print(SocialNetwork['TWITTER'])
SocialNetwork.TWITTER
>>> print(SocialNetwork(1))
SocialNetwork.TWITTER

Each member has dedicated attributes for their name and value:

>>> print(SocialNetwork.TIKTOK.value)
2
>>> print(SocialNetwork(2).name)
TIKTOK

You can use the list() function to get all the members of an enumeration, including their names and values:

>>> print(list(SocialNetwork))
[<SocialNetwork.TWITTER: 1>, <SocialNetwork.TIKTOK: 2>, <SocialNetwork.FACEBOOK: 3>, <SocialNetwork.INSTAGRAM: 4>]

The isinstance() function can help you check whether a name is a member of an enumeration.

>>> isinstance(SocialNetwork.FACEBOOK, SocialNetwork)
True

Loop Through An Enumeration

One of the main key properties of enumerations is that they are iterable, meaning you can loop over an enumeration and get its members in order.

>>> for i in SocialNetwork:
...         print(i)
... 
SocialNetwork.TWITTER
SocialNetwork.TIKTOK
SocialNetwork.FACEBOOK
SocialNetwork.INSTAGRAM

Compare Enumerations

You can compare members of an enumeration by using equality operators:

>>> SocialNetwork(1) == SocialNetwork.TWITTER
True
>>> SocialNetwork.TIKTOK != SocialNetwork.TWITTER
True

The is keyword also works with enumeration members:

>>> SocialNetwork(1) is SocialNetwork.TWITTER
True
>>> SocialNetwork.TIKTOK is not SocialNetwork.TWITTER
True

Remember that you can use comparison operators like > or < because enumeration doesn’t support ordered comparison by default:

>>> SocialNetwork.TIKTOK > SocialNetwork.TWITTER
TypeError: '>' not supported between instances of 'SocialNetwork' and 'SocialNetwork'

To use those operators, create your enumeration from the IntEnum class:

from enum import IntEnum
class Order(IntEnum):
        FIRST = 1
        SECOND = 2
        THIRD = 3
        FOURTH = 4

You will be able to use its members anywhere integers can be used:

>>> Order.FIRST < Order.SECOND
True
>>> Order.FIRST + Order.SECOND
3
>>> Order.FIRST == 1
True

Additionally, Python 3.11 adds the StrEnum class:

from enum import StrEnum
class Domain(StrEnum):
        FIRST = 'learnshareit'
        SECOND = '.'
        THIRD = 'com'

 Members of this enumeration will behave like strings:

>>> Domain.FIRST + Domain.SECOND + Domain.THIRD
 'learnshareit.com'

Add Methods To An Enumeration

Since Python numerations are classes, you can extend and customize their behaviors by adding class and instance methods.

In this example, we define an enumeration of four simple stages in software development, from design to release. We then create an instance method that returns the next member in the enumeration

class ProductStatus(Enum):
        DESIGN = 1
        DEVELOP = 2
        TEST = 3
        RELEASE = 4
        def next(self):
                cls = self.__class__
                return cls(self.value + 1)

Here is how you can use the next() method:

>>> status = ProductStatus(1)
>>> print(status)
ProductStatus.DESIGN
>>> status = status.next()
>>> print(status)
ProductStatus.DEVELOP

You can also create class methods with the @classmethod decorator:

@classmethod
def total(cls):
        n = 0
        for i in cls:
                n = n + i.value
        return n

This class method returns the sum of all the values of members in the enumerations:   

 >>> print(ProductStatus.total())
 10

Tutorials

You can learn more about Enum in Python in the articles below:

Summary

The Python enum module provides support for enumerations with the Enum class and other derivations like IntEnum or StrEnum. They are useful when you need to map integers (or other values) to a fixed set of possible values.

Related Posts

Leave a Reply

Your email address will not be published. Required fields are marked *