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)
isstr
and notUSER_TYPES
-
user_a.type == "Admin"
isTrue
- 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
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
andCUSTOMER
are0
, so it's evaluated tofalsy
-
type(user_a.type)
isint
and notUserType
-
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 compareUserType.CUSTOMER
to other ints -
UserType.CUSTOMER
is falsy, because its value is0
- 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
- PEP 435
- Python Enum
- How are Enums different?
- Enumerated type
- SQL Alchemy Enum
- Django Enumeration Types
Top comments (0)