简体   繁体   中英

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?

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

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

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM