DEV Community

Christian Barra
Christian Barra

Posted on • 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
Enter fullscreen mode Exit fullscreen mode

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
Enter fullscreen mode Exit fullscreen mode

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 constants1:

from dataclasses import dataclass

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

@dataclass
class User:
    name: str
    type: int
Enter fullscreen mode Exit fullscreen mode 
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
Enter fullscreen mode Exit fullscreen mode

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
Enter fullscreen mode Exit fullscreen mode

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
Enter fullscreen mode Exit fullscreen mode

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
Enter fullscreen mode Exit fullscreen mode

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
Enter fullscreen mode Exit fullscreen mode 
# 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>]
Enter fullscreen mode Exit fullscreen mode

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>
Enter fullscreen mode Exit fullscreen mode

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

UserType.ADMIN.value
# 2

UserType.ADMIN.name
# ADMIN
Enter fullscreen mode Exit fullscreen mode

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
Enter fullscreen mode Exit fullscreen mode

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

  1. Python constants ↩

Top comments (0)

Read next

jamesli profile image

Building Enterprise Agent Systems: Core Component Design and Optimization

James Li -

yassinesallami profile image

#? List vs Tuples in python

Yassine Sallami -

brandonc profile image

DUCK (file structure) YOU!

Brandon Collins -

lakshmi_prithanadesan_32 profile image

Day 12 - Looping Excercises

Lakshmi Pritha Nadesan -