How to Use Enums with Pydantic (3 Examples)

Updated: November 30, 2023 By: Khue Post a comment

This tutorial shows you how to use enums with Pydantic.

Overview

Pydantic is a library that helps you validate and parse data using Python type annotations. Enums help you create a set of named constants that can be used as valid values for a variable or a parameter. Pydantic supports Python’s standard enum classes to define choices.

In general, the process of using enums with Pydantic looks as follows:

  1.  import the Enum class from the enum module and the BaseModel class from the pydantic module.
  2. Define your own enum class by subclassing the Enum class and assigning the possible values as class attributes. Optionally, you can also utilize the subclass str or int to make your enum values of that type.
  3. Define your Pydantic model by subclassing the BaseModel class and annotating the fields with your enum class or a specific member of it. You can also provide a default value for the fields.
  4.  Initialize an instance of your Pydantic model by passing the enum values or instances as arguments or keyword arguments. Pydantic will validate and parse the data according to your enum definition.
  5. You can access the fields of your Pydantic model instance as attributes and get their values or names using the .value or .name properties.

It’s tough to learn Pydantic with just words. It’s time to write some actual code to get a better and deeper understanding.

Examples

Basic Enum Usage

This example shows how to define an enum class and use it as a type annotation for a Pydantic model field. Pydantic will validate that the field value is a valid member of the enum class.

from enum import Enum
from pydantic import BaseModel, ValidationError

# Define an Enum for the color of cars
class Color(str, Enum):
    red = "red"
    green = "green"
    blue = "blue"

# Define a Pydantic model that presents a car
class Car(BaseModel):
    make: str
    color: Color

# Create a car with a valid color
car1 = Car(make="Toyota", color=Color.red)
print(car1)
# Output: make='Toyota' color=<Color.red: 'red'>

# Create another car with an valid color
car2 = Car(make="Tesla", color="blue")
print(car2)
# Output: make='Tesla' color=<Color.blue: 'blue'>

If you try to create a car object with an invalid color, the ValidationError will be raised:

# Create a car with an invalid color
car3 = Car(make="Ford", color="yellow")
print(car3)
# ValidationError: 1 validation error for Car
# color
# Input should be 'red','green' or 'blue' [type=enum, input_value='yellow', input_type=str]

Literal enum usage

This example shows how to use the Literal type from the typing module to restrict a Pydantic model field to one specific enum value. It also shows how to create instances of the model using both the enum instance and the enum value.

from enum import Enum
from typing import Literal
from pydantic import BaseModel, ValidationError

# Define enum
class SizeEnum(str, Enum):
    BIG = 'big'
    AVERAGE = 'average'
    SMALL = 'small'

# Define model
class MyModel(BaseModel):
    x: Literal[SizeEnum.AVERAGE, SizeEnum.SMALL] # Only accept average or small

print(MyModel(x=SizeEnum.AVERAGE)) # x=<SizeEnum.AVERAGE: 'average'>
print(MyModel(x='average')) # x=<SizeEnum.AVERAGE: 'average'>
print(MyModel(x=SizeEnum.SMALL)) # x=<SizeEnum.SMALL: 'small'>

try:
    MyModel(x=SizeEnum.BIG) # invalid value
except ValidationError as e:
    print(e) 
    # 1 validation error for MyModel
    # x
    #   Input should be <SizeEnum.AVERAGE: 'average'> or <SizeEnum.SMALL: 'small'> [type=literal_error, input_value=<SizeEnum.BIG: 'big'>, input_type=SizeEnum]

List of enums usage

This example shows how to use a list of enums as a type for a Pydantic model field. It also demonstrates how to use enum values instead of keys when printing objects initialized with the Pydantic model (with the config use_enum_values = True).

from pydantic import BaseModel
from typing import List
from enum import Enum

# Define Hobby Enum
class Hobby(str, Enum):
    SWIMMING = "swimming"
    READING = "reading"
    SINGING = "singing"

# Define Person Model
class Person(BaseModel):
    name: str
    hobbies: List[Hobby] = []

    class Config:
        # Use Enum values instead of keys
        use_enum_values = True
    
# Create some instances of Person
p1 = Person(name="John", hobbies=[Hobby.SWIMMING, Hobby.READING])
print(p1.hobbies) # ['swimming', 'reading']

p2 = Person(name="Mary", hobbies=["singing"])
print(p2.hobbies) # ['singing']

Conclusion

You’ve learned the fundamentals of using enumerations with Pydantic models. You’ve also seen several examples (from basic to advanced) of applying that knowledge in practice. From this point, you can write more robust and safe programs. This tutorial ends here. Happy coding & enjoy your day!