Christian Barra

Posted on Dec 14, 2020 • Originally published at pybootcamp.com

Why you need enums and how to use them in Python

#python

Let's say that you have a user object with two attributes: name and type.

A user's type can be Admin, Moderator and Customer.

How would you model it?

Using a Tuple

One way is to use a tuple to hold type's values:

from dataclasses import dataclass

USER_TYPES = (
    "Customer",
    "Moderator",
    "Admin",
)

@dataclass
class User:
    name: str
    type: str

Let's try and see if a tuple can do the job:

# Admin
user_a = User("Christian", USER_TYPES[2])

# Customer
user_b = User("Daniel", USER_TYPES[0])

user_a.type == user_b.type
# False

# Ideally it should be USER_TYPES 😞
type(user_a.type)
# <class 'str'>

# Ideally this should be False 😞
user_a.type == "Admin"
# True

len(USER_TYPES)
# 3

[user_type for user_type in USER_TYPES]
# ['Customer', 'Moderator', 'Admin']

if user_b.type:
    print("User is valid")
else:
    print("Not a valid user")
# User is valid

It's a good first attempt, but there are few caveats that I want to address:

type(user_a.type) is str and not USER_TYPES
user_a.type == "Admin" is True
who's gonna remember what USER_TYPES[2] is?

Addressing the last point is easy, but the first two are tricky: I need to the separate the UserType from its concrete representation (in this case a string).

Using Constants

Another approach to consider is to assign numerical values to constants¹:

from dataclasses import dataclass

CUSTOMER, MODERATOR, ADMIN = range(0,3)

@dataclass
class User:
    name: str
    type: int

user_a = User("Christian", ADMIN)
user_b = User("Daniel", CUSTOMER)

user_a.type == user_b.type
# False

type(user_1.type)
# <class 'int'>

user_a.type == "admin"
# False

# Ideally this should be False 😞
user_b.type == 0
# True

if user_b.type == CUSTOMER:
    print("User type: CUSTOMER")
# User type: CUSTOMER

# `user_b.type` should evaluate to True 😞
if user_b.type:
    print("User is valid")
else:
    print("Not a valid user")
# Not a valid user

Using constants really improves readability: I can now use CUSTOMER, MODERATOR and ADMIN.

But it also has some (new) drawbacks:

user_b.type and CUSTOMER are 0, so it's evaluated to falsy
type(user_a.type) is int and not UserType
user_b.type == 0 is True
len(USER_TYPES) is gone, I don't have a way to reference to the collection of user's types

Using a Class

Using a class UserType can be a way to unify my first two attempts.

from dataclasses import dataclass

class UserType:
    CUSTOMER = 0
    MODERATOR = 1
    ADMIN = 2

@dataclass
class User:
    name: str
    # NOTE: type is not UserType 😞
    type: int

Let's see how the UserType behaves:

user_a = User(name="Christian", type=UserType.ADMIN)
user_b = User(name="Daniel", type=UserType.CUSTOMER)

user_a.type == user_b.type
# False

type(user_a.type)
# <class 'int'>

user_a.type == "ADMIN"
# False

user_b.type == 0
# True

len(UserType)
# Traceback (most recent call last):
#   File "<stdin>", line 1, in <module>
# TypeError: object of type 'type' has no len()

if user_b.type:
    print("User is valid")
else:
    print("Not a valid user")
# Not a valid user 😞

[user_type for user_type in UserType]
# Traceback (most recent call last):
#   File "<stdin>", line 1, in <module>
# TypeError: 'type' object is not iterable

I grouped values together, under UserType, and it's definitely more readable than before.

But there are still a few drawbacks, similar to the tuple implementation:

the underlying concrete type for each value is int so you can compare UserType.CUSTOMER to other ints
UserType.CUSTOMER is falsy, because its value is 0
classes are mutable, and the fact that I can modify the class at runtime isn't exactly what I want

Having a class that encapsulates the concept of type/category is useful, but a simple class is leaving us half way.

I need a (new) type that supports enumeration, something that I can count, ideally immutable.

And I need a (new) type where the underlying representation is somehow abstracted away, such as:

type(UserType.ADMIN) is str
# False

type(UserType.ADMIN) is int
# False

The type that I need is called Enum and Python supports Enum, they were introduced with PEP 435 and they are part of the standard library.

From PEP 435:

An enumeration is a set of symbolic names bound to unique, constant values.
Within an enumeration, the values can be compared by identity, and the enumeration itself can be iterated over.

Using Enums

Enum types are data types that comprise a static, ordered set of values.

An example of an enum type might be the days of the week, or a set of status values for a piece of data (like my User's type).

from dataclasses import dataclass
from enum import Enum

class UserType(Enum):
    CUSTOMER = 0
    MODERATOR = 1
    ADMIN = 2

@dataclass
class User:
    name: str
    type: UserType

# exactly like with the class example
user_a = User(name="Christian", type=UserType.ADMIN)
user_b = User(name="Daniel", type=UserType.CUSTOMER)

user_a.type == user_b.type
# False

type(user_a.type)
# <enum 'UserType'> 🚀

user_a.type == "ADMIN"
# False

user_a.type == 2
# False

len(UserType)
# 3

if user_b.type:
    print("User is valid")
else:
    print("Not a valid user")
# User is valid

[user_type for user_type in UserType]
# [<UserType.CUSTOMER: 0>, <UserType.MODERATOR: 1>, <UserType.ADMIN: 2>]

So an Enum is already more powerful than a simple class in modelling our business case. It really helps us express some of the characteristics of a category/list.

There are a couple of more things related to enums, some nice built-in features.

I can access an Enum in different ways, using both brackets and dot notations:

UserType['ADMIN']
# <UserType.ADMIN: 2>

UserType.ADMIN
# <UserType.ADMIN: 2>

Each member of the Enum has a value and name attribute:

UserType.ADMIN.value
# 2

UserType.ADMIN.name
# ADMIN

Another neat feature is that you can't really modify an Enum at runtime:

list(UserType)
# [<UserType.CUSTOMER: 0>, <UserType.MODERATOR: 1>, <UserType.ADMIN: 2>]

UserType.NOOB = 3

list(UserType)
# [<UserType.CUSTOMER: 0>, <UserType.MODERATOR: 1>, <UserType.ADMIN: 2>]

type(UserType.ADMIN) is type(UserType.MODERATOR)
# True

type(UserType.ADMIN) is type(UserType.NOOB)
# False

As you can see enums in Python are really powerful, and they encapsulate well concepts like categories or types.

The Python documentation is great and thoroughly covers all the details, I'd recommend checking it out.

Useful resources

Python constants ↩

DEV Community

Why you need enums and how to use them in Python

Using a Tuple

Using Constants

Using a Class

Using Enums

Useful resources

Top comments (0)

Read next

Performance trap: general libraries & helper objects

Concurrency and Parallelism in Python

A look into the new module in Python 3.13.0 - dbm.sqlite3

🎉 Dive into Python 3.13: New Features, JIT Compiler, Free threading & Installation Guide! 🎉