encode Pydantic field using the Enum name instead of the value
Asked Answered
C

7

10

I have a Enum class:

class Group(enum.Enum):
    user = 0
    manager = 1
    admin = 2

I have a pydantic model:

class User(BaseModel):
    id: int
    username: str
    group: Group

It generated serialised to json following:

{
    "id": 5,
    "username": "admin",
    "group": 2
}

However, I would like to get the name of the enum field instead of its value, so it should be:

{
    "id": 5,
    "username": "admin",
    "group": "admin"
}

Is this possible at all? If so, how?

Caracas answered 9/6, 2021 at 20:32 Comment(1)
For clues on serializing/deserializing via JSON start here: https://mcmap.net/q/183588/-is-it-possible-to-dump-an-enum-in-json-without-passing-an-encoder-to-json-dumpsShortage
S
4

You are not trying to "load" anything, your problem is that you want to encode your Pydantic field using the Enum name instead of the value, when serialising your model to JSON.

All you need to do is add a Config class to your BaseModel subclass that specifies a JSON encoder for the Group type. The json_encoders attribute is a dict keyed on type with serialiser callables:

import enum
from pydantic import BaseModel


class Group(enum.Enum):
    user = 0
    manager = 1
    admin = 2


class User(BaseModel):
    id: int
    username: str
    group: Group

    class Config:
        json_encoders = {Group: lambda g: g.name}


user = User(id=5, username="admin", group=2)
print(user)  # id=5 username='admin' group=<Group.admin: 2>
print(user.json())  # {"id": 5, "username": "admin", "group": "admin"}
Sannyasi answered 27/2, 2023 at 23:20 Comment(2)
For others who may want to go further, and actually validate their Enum by name, that is to create the BaseModel using Enum names instead of values, try this: github.com/pydantic/pydantic/discussions/2980Sannyasi
There is also this question that tackles customised Enums including modifying the schema and validating input by name instead of value - and the Enum is loaded dynamically: #75587942Sannyasi
F
8

The method given in the accepted answer has been deprecated for Pydantic V2. The V2 method is to use custom serializer decorators, so the way to do this would now look like so:

import enum
from pydantic import BaseModel, field_serializer


class Group(enum.Enum):
    user = 0
    manager = 1
    admin = 2


class User(BaseModel):
    id: int
    username: str
    group: Group

    @field_serializer("group")
    def serialize_group(self, group: Group, _info):
        return group.name
Falsework answered 6/10, 2023 at 22:29 Comment(3)
Nice. Many of the other answers apply the enum names instead of values for all fields in the model (or all invocations of the enum). This method allows per-field control, such that a model could have one field that uses the enum values and another that uses the enum names.Testes
I just tried this approach to generate a clean OpenAPI documentation with FastAPI for a route response model... and it does not work (yet?) Pydantic v2.4.2, FastAPI 0.103.0.Glengarry
Maybe there is something in the way FastAPI does the JSON conversion that doesn't quite work. I use the model_dump_json method of BaseModel directly and have had no problems (Pydantic v2.4.1).Falsework
S
4

You are not trying to "load" anything, your problem is that you want to encode your Pydantic field using the Enum name instead of the value, when serialising your model to JSON.

All you need to do is add a Config class to your BaseModel subclass that specifies a JSON encoder for the Group type. The json_encoders attribute is a dict keyed on type with serialiser callables:

import enum
from pydantic import BaseModel


class Group(enum.Enum):
    user = 0
    manager = 1
    admin = 2


class User(BaseModel):
    id: int
    username: str
    group: Group

    class Config:
        json_encoders = {Group: lambda g: g.name}


user = User(id=5, username="admin", group=2)
print(user)  # id=5 username='admin' group=<Group.admin: 2>
print(user.json())  # {"id": 5, "username": "admin", "group": "admin"}
Sannyasi answered 27/2, 2023 at 23:20 Comment(2)
For others who may want to go further, and actually validate their Enum by name, that is to create the BaseModel using Enum names instead of values, try this: github.com/pydantic/pydantic/discussions/2980Sannyasi
There is also this question that tackles customised Enums including modifying the schema and validating input by name instead of value - and the Enum is loaded dynamically: #75587942Sannyasi
C
3

In 2.5 this is simple using the PlainSerializer.

Relevant Docs

import pydantic, enum
from typing import Annotated

EnumNameSerializer = pydantic.PlainSerializer(
    lambda e: e.name,
    return_type='str',
    when_used='always')

class Group(enum.Enum):
    user = 0
    manager = 1
    admin = 2

class User(pydantic.BaseModel):
    id: int
    username: str
    group: Annotated[Group, EnumNameSerializer]

print(User(group=Group.admin, id=0, username='admin_user').model_dump_json())

>>> {"id":0,"username":"admin_user","group":"admin"}

Note the change of type for the group attribute:

group: Annotated[Group, EnumNameSerializer]

Since this uses when_used='always', this will also output the enum name when exporting to a dict with .model_dump()

Clancy answered 9/12, 2023 at 0:31 Comment(1)
Is there a way of doing this so that User.model_validate_json(json_data=some_user.model_dump_json()) Works?Aussie
C
1

You could override the dunder method repr in the Enum class in this way:

class Group(enum.Enum):
    user = 0
    manager = 1
    admin = 2

    def __repr__(self) -> str:
        return self.name
Clash answered 12/5, 2023 at 6:5 Comment(1)
Answer needs supporting information Your answer could be improved with additional supporting information. Please edit to add further details, such as citations or documentation, so that others can confirm that your answer is correct. You can find more information on how to write good answers in the help center.Strip
M
0

In the pydantic Github I found a discussion tackling this issue: Map string value to int #598

samuelcolvin suggested this:

class Choices(int, Enum):
    anne = 1
    ben = 2
    charlie = 3
    dave = 4

    @classmethod
    def __get_validators__(cls):
        cls.lookup = {v: k.value for v, k in cls.__members__.items()}
        yield cls.validate

    @classmethod
    def validate(cls, v):
        try:
            return cls.lookup[v]
        except KeyError:
            raise ValueError('invalid value')

class Model(BaseModel):
    choice: Choices

debug(Model(choice='charlie'))
Magnificence answered 1/2, 2023 at 7:53 Comment(2)
this doesnt work.Kentkenta
It works with pydantic==1.10.6, I tested it. Maybe there are problems with pydantic v2?Magnificence
G
0

Not answering the question directly, but it's related.

If you want to bind an enum to a pydantic model without relying on its value, you can create an enum on the fly. (Python >= 3.11) This works with FastAPI and the generated OpenAPI schema will reflect it properly.

class Group(Enum):
    user = 0
    manager = 1
    admin = 2


group_enum_names = Enum("GroupEnumNames", [(g.name, g.name) for g in Group])


class User(BaseModel):
    id: int
    username: str
    group: group_enum_names


# works
print(User(**{"id": 5, "username": "admin", "group": "admin"}))
# yields {"id":5,"username":"admin","group":"admin"}

# does not work, there are no ties to the original enum anymore
print(User(**{"id": 5, "username": "admin", "group": 1}))
# yields pydantic_core._pydantic_core.ValidationError

And to map back to the original Group enum:

user = User(**{"id": 5, "username": "admin", "group": "admin"})
group = Group[user.group.name]

Arguably, some logic to could be added directly into the User model to handle that.

Note: mypy 1.6.1 does not like that solution and yields Second argument of Enum() must be string, tuple, list or dict literal for mypy to determine Enum members [misc] over the dynamically created enum, even if the second argument is a tuple.

Tested with:

  • pydantic v2.4.2
  • fastapi v0.103.0
Glengarry answered 20/10, 2023 at 1:17 Comment(0)
C
-1
class Group(enum.Enum):
    user = "user"
    manager = "manager"
    admin = "admin"

did the trick :o 🎉

Caracas answered 9/6, 2021 at 20:51 Comment(3)
Do you need to deserialize that back from JSON to Python? How would you do it?Shortage
Hopefully the OP can add a complete example showing pydantic both serializing and unserializing some data.Shortage
No, its done by fastapi ... I do not serialize by myself.Caracas

© 2022 - 2024 — McMap. All rights reserved.