CODEWITHSUNDEEP
pythonjsonserializationenumspydantic
RaamEE
RaamEE

Reputation: 3507

How to JSON serialize ENum classes in Pydantic BaseModel

I have the following code that uses Pydantic BaseModel data class

from enum import Enum

import requests
from pydantic import BaseModel
from requests import Response


class PetType(Enum):
    DOG: str = 'dog'
    CAT: str = 'cat'


class Pet(BaseModel):
    name: str
    type: PetType


my_dog: Pet = Pet(name='Lucky', type=PetType.DOG)

# This works
resp: Response = requests.post('https://postman-echo.com/post', json=my_dog.json())
print(resp.json())

#This doesn't work
resp: Response = requests.post('https://postman-echo.com/post', json=my_dog.dict())
print(resp.json())

That when I send json equals to model's dict(), I get the error:

TypeError: Object of type 'PetType' is not JSON serializable

How do I overcome this error and make PetType also serializable?

P.S. The above example is short and simple, but I hit a use case where both cases of sending

json=my_dog.json()

and

json=my_dog.dict()

don't work. This is why I need to solve sending using dict().

Upvotes: 9

Views: 16407

Answers (2)

cottontail
cottontail

Reputation: 23261

In Pydantic 2, with the models defined exactly as in the OP, when creating a dictionary using model_dump, we can pass mode="json" to ensure that the output will only contain JSON serializable types. Then, working off of the code in the OP, we could change the post request as follows to get the desired behavior:

di = my_dog.model_dump(mode="json")    # <--- convert to serializable dict
resp = requests.post('https://postman-echo.com/post', json=di)
print(resp.json())

Yet another way is to pass use_enum_values=True in the model definition. This is because normally, Enum values can be obtained by .value attribute and in Pydantic BaseModels, that can be done by passing use_enum_values as True in the object definition.

class Pet(BaseModel, use_enum_values=True):   # <--- specify here
    name: str
    type: PetType

my_dog = Pet(name='Lucky', type=PetType.DOG)

resp = requests.post('https://postman-echo.com/post', data=my_dog.model_dump_json())
print(resp.json())

resp = requests.post('https://postman-echo.com/post', json=my_dog.model_dump())
print(resp.json())

Upvotes: 0

RaamEE
RaamEE

Reputation: 3507

**<---- Addition 2 ----> **

Check types like https://docs.python.org/3/library/enum.html#enum.StrEnum and https://docs.python.org/3.12/library/enum.html#enum.IntEnum

Instead of MyEnum(str, Enum) use MyEnum(StrENum)

**<---- Addition ----> **

Look for Pydantic's parameter "use_enum_values" in Pydantic Model Config

use_enum_values whether to populate models with the value property of enums, rather than the raw enum. This may be useful if you want to serialise model.dict() later (default: False)

It looks like setting this value to True will do the same as the below solution.

Turns out that this is a behavior of ENum, which is discussed here: https://github.com/samuelcolvin/pydantic/issues/2278

The way you should define the enum is using

class PetType(str, Enum):

instead of

class PetType(Enum):

For integers this Python's Enum library provides the type IntEnum: https://docs.python.org/3.10/library/enum.html#enum.IntEnum

which is basically

class IntEnum(int, Enum):
    pass

If you look at the above Enum documentation you will find that a type like StrEnum doesn't exist but following the example for PetType you can define it easily.

I am attaching the working code below

from enum import Enum

import requests
from pydantic import BaseModel
from requests import Response


class PetType(str, Enum):
    DOG: str = 'dog'
    CAT: str = 'cat'


class Pet(BaseModel):
    name: str
    type: PetType


my_dog: Pet = Pet(name='Lucky', type=PetType.DOG)

# This works
resp: Response = requests.post('https://postman-echo.com/post', json=my_dog.json())
print(resp.json())

# Now this also works
resp: Response = requests.post('https://postman-echo.com/post', json=my_dog.dict())
print(resp.json())

Upvotes: 18

Related Questions