73

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

phoenix
  • 7,988
  • 6
  • 39
  • 45
ierdna
  • 5,753
  • 7
  • 50
  • 84

4 Answers4

141

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())
alex_noname
  • 26,459
  • 5
  • 69
  • 86
  • 24
    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. – Krishna Jul 21 '21 at 10:30
  • 7
    @Krishna 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 Enum – Stephen Carboni Sep 19 '21 at 15:09
  • 8
    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 ? – vianmixt Nov 18 '21 at 09:45
  • 4
    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. – masus04 Apr 08 '22 at 12:21
  • @vianmixt See [y26805's answer](https://stackoverflow.com/a/73697224) for a solution that relies on FastAPI. – phoenix Oct 24 '22 at 19:15
7

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'}
phoenix
  • 7,988
  • 6
  • 39
  • 45
y26805
  • 71
  • 1
  • 3
  • This is nice because you can keep the full `Enum` when deserializing and constructing the Pydantic object. – phoenix Oct 24 '22 at 19:14
3

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.

mr_mo
  • 1,401
  • 6
  • 10
1

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
Yaakov Bressler
  • 9,056
  • 2
  • 45
  • 69