2022-03-30

Enum in Python

What is Enum

An enumeration, often shortened to enum, is a distinct data type that consists of a set of named values called elements or members. In other words, an enum is a way of grouping related constants together, thereby creating a new type that can take on one of a few predefined values.

For example, consider a scenario where we are developing a system that needs to keep track of different statuses for a task. These statuses could be "Not Started", "In Progress", "Completed", and "On Hold". Here, we could create an enum named TaskStatus with these four statuses as the members of the enum. This gives us a neat, organized, and error-resistant way to handle task statuses in our system.

Enum in Python

We will explore how to use enumerations, or enums, in Python. We will cover how to define an enum, how to access its members, and how to iterate over an enum.

Enum Module

Python provides the enum module for creating enumerations. This module defines a base class Enum which can be used to derive any enumeration.

python
from enum import Enum

Defining Enum

To define an enumeration in Python, you create a class that inherits from Enum and add class attributes for each possible value.

For example, we can define an enumeration Day for days of the week:

python
from enum import Enum

class Day(Enum):
    MONDAY = 1
    TUESDAY = 2
    WEDNESDAY = 3
    THURSDAY = 4
    FRIDAY = 5
    SATURDAY = 6
    SUNDAY = 7

In this example, Day is the enum, and MONDAY, TUESDAY, WEDNESDAY, etc., are its members. Each member has a name (e.g., MONDAY) and a value (e.g., 1).

Accessing Enum Members

After defining our Day enumeration, we can access its members directly using dot notation:

python
print(Day.MONDAY)
print(Day.TUESDAY)
Day.MONDAY
Day.TUESDAY

We can also access the value of an enumeration member using the .value property:

python
print(Day.MONDAY.value)

This will output: 1.

Iterating through Enum

Enumerations are iterable. They can be looped over using a for loop:

python
for day in Day:
    print(day)
Day.MONDAY
Day.TUESDAY
Day.WEDNESDAY
Day.THURSDAY
Day.FRIDAY
Day.SATURDAY
Day.SUNDAY

Comparing Enums

Enum members can be compared using their identity or their value. When you compare enum members using their identity (i.e., using is), you're comparing whether they are the exact same member of the same enum. When you compare enum members using their value (i.e., using ==), you're comparing whether they have the same value.

Here's an example:

python
print(Day.MONDAY is Day.MONDAY) # True
print(Day.MONDAY is Day.TUESDAY) # False
print(Day.MONDAY == Day.MONDAY) # True
print(Day.MONDAY == Day.TUESDAY) # False
print(Day.MONDAY == 1) # True

Aliases in Enum

In Python enums, you can have two members with the same value. When this happens, the second member is considered an alias of the first member. Here's an example:

python
from enum import Enum

class Color(Enum):
    RED = 1
    BLUE = 2
    GREEN = 3
    YELLOW = 4
    ORANGE = 1

In this case, ORANGE is an alias for RED. If you try to access Color.ORANGE, you'll get Color.RED instead:

python
print(Color.ORANGE)  # Output: Color.RED

Auto() in Enum

If you don't care about the actual value associated with each member of your enum, you can use the auto() function to automatically assign values. Here's an example:

python
from enum import Enum, auto

class Color(Enum):
    RED = auto()
    BLUE = auto()
    GREEN = auto()
    YELLOW = auto()

In this case, Python will automatically assign the values 1, 2, 3, and 4 to RED, BLUE, GREEN, and YELLOW respectively.

Ensuring Unique Enum Members

If you want to ensure that all members of your enum have unique values, you can use the @unique decorator from the enum module. If you apply this decorator to an enum with non-unique values, Python will raise a ValueError.

python
from enum import Enum, unique

@unique
class Color(Enum):
    RED = 1
    BLUE = 2
    GREEN = 3
    YELLOW = 4
    ORANGE = 5

In this case, if you try to add another color with a value that's already used (say, PINK = 1), Python will raise a ValueError.

Ryusei Kakujo

researchgatelinkedingithub

Focusing on data science for mobility

Bench Press 100kg!