Pydantic enum field does not get converted to string
Asked Answered
A

4

89

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'}

Antimicrobial answered 9/12, 2020 at 2:50 Comment(0)
C
161

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())
Carbazole answered 9/12, 2020 at 6:32 Comment(6)
is there a reason that this option is False by default? I feel in most of the cases we would require that to be true.Salonika
@Salonika one could argue why not leave dates as ISO8601 strings? Principle of least surprise is that you get what you ask for, you asked for an EnumCovering
Is there a way to use the enum values when serializing using .dict() but keep the full enum when deserializing and constructing the pydantic object ?Champlain
I was wondering the same as @vianmixt. The whole reason why anyone would use an enum is so that when working with the object, one has the convenience that enums offer, but yet it can be serialized as a string as this is the simplest way to be language independent.Rothwell
@Champlain See y26805's answer for a solution that relies on FastAPI.Danseuse
Actually when using straight str based Enum where each member has the same name and value (e.g. MALE = "MALE" ; FEMALE = "FEMALE" etc.) I don't find any convenience in having the Enum instance. It's more verbose and requires a reference. Given Pydantic has validated the input nicely I can rely on a derived str being one of the allowed values and can use less verbose string comparisons if I need to. The primary reason I have for using Enum is to create the OpenAPI description.Corking
U
14

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
Uredium answered 24/8, 2023 at 23:11 Comment(2)
Thanks for the tip about the kwargss!Crashland
Calling .model_dump(mode="json") allows for being able to get a json-able dict when dumping the object while retaining the ability to construct and deserialize objects with the actual enum type.Nodababus
F
11

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'}
Fiona answered 13/9, 2022 at 3:21 Comment(2)
This is nice because you can keep the full Enum when deserializing and constructing the Pydantic object.Danseuse
This is my preferred way of doing this as you don't always want to loose the enum type when using model_dump, would be nicer to pass use_enum_values when calling model_dump. Using model_dump_json does not always work in our use case either.Pivoting
M
9

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.

Munificent answered 10/5, 2023 at 23:6 Comment(0)

© 2022 - 2024 — McMap. All rights reserved.