CODEWITHSUNDEEP
pythonserializationfastapipydantic
ierdna
ierdna

Reputation: 6293

Pydantic enum field does not get converted to string

I am trying to restrict one field in a class to an enum. However, when I try to get a dictionary out of class, it doesn't get converted to string. Instead it retains the enum. I checked pydantic documentation, but couldn't find anything relevant to my problem.

This code is representative of what I actually need.

from enum import Enum
from pydantic import BaseModel

class S(str, Enum):
    am = 'am'
    pm = 'pm'

class K(BaseModel):
    k: S
    z: str

a = K(k='am', z='rrrr')
print(a.dict()) # {'k': <S.am: 'am'>, 'z': 'rrrr'}

I'm trying to get the .dict() method to return {'k': 'am', 'z': 'rrrr'}

Upvotes: 101

Views: 95199

Answers (4)

Yaakov Bressler
Yaakov Bressler

Reputation: 12088

Pydantic 2.0

By default, Pydantic preserves the enum data type in its serialization. To override this behavior, specify use_enum_values in the model config.

from enum import Enum
from pydantic import BaseModel, ConfigDict

class S(str, Enum):
    am = 'am'
    pm = 'pm'


class K(BaseModel):
    model_config = ConfigDict(use_enum_values=True)

    k: S
    z: str

a = K(k='am', z='rrrr')
print(a.model_dump()) # {'k': 'am', 'z': 'rrrr'}

Note: model_config is now an attribute of type ConfigDict (this is a breaking change from V1).

Neater: Specify config options as model class kwargs

Alternatively, you can populate directly in the class definition:

class K(BaseModel, use_enum_values=True):
    k: S
    z: str

Upvotes: 26

mr_mo
mr_mo

Reputation: 1528

The answer by using Config class is correct but may not be necessary. It may be inconvenient to repeatedly implement and could be redundant for some readers so I offer easier alternatives and a short explanation.

TL;DR

  • For JSON serialization, no need to do anything. Using (Str,Enum) subclasses is enough without Config.
  • For printable representational string/other feature, create a custom StrEnum class. This was added in 3.11 but since previous versions don't support it - I provide an example.

Example

from enum import Enum

class StrEnum(str, Enum):
    def __repr__(self) -> str:
        return str.__repr__(self.value)

class A(str, Enum):
    FOO = "foo"

class B(StrEnum):
    BAR= "bar"

class C(BaseModel):
    a: A = Field(...)
    b: B = Field(...)

print(C(a="foo", b="bar").dict())

import json

print(json.dumps(C(a="foo", b="bar").dict()))

Outputs:

{'a': <A.FOO: 'foo'>, 'b': 'bar'}
{"a": "foo", "b": "bar"}

Explanation

Class A and B are of types (str, Enum) and StrEnum, respectively. Both are JSON serializable; StrEnum also prints as a primitive string.

Upvotes: 9

y26805
y26805

Reputation: 131

You can use FastAPI's jsonable_encoder:

from enum import Enum
from pydantic import BaseModel
from fastapi.encoders import jsonable_encoder

class S(str, Enum):
    am = 'am'
    pm = 'pm'

class K(BaseModel):
    k: S
    z: str

a = K(k='am', z='rrrr')
print(jsonable_encoder(a)) # {'k': 'am', 'z': 'rrrr'}

Upvotes: 13

alex_noname
alex_noname

Reputation: 32233

You need to use use_enum_values option of 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)

from enum import Enum
from pydantic import BaseModel

class S(str, Enum):
    am='am'
    pm='pm'

class K(BaseModel):
    k:S
    z:str

    class Config:  
        use_enum_values = True  # <--

a = K(k='am', z='rrrr')
print(a.dict())

Upvotes: 182

Related Questions