如何 JSON 序列化 Pydantic BaseModel 中的 ENum 类

Question

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?

PS 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()

Answer 1

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())