使用枚举名称而不是值对 Pydantic 字段进行编码

Question

我有一个枚举类：

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

我有一个 pydantic 模型：

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

它生成序列化为 json 的内容如下：

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

但是，我想获取枚举字段的名称而不是其值，所以它应该是：

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

这有可能吗？如果是这样，怎么办？

Answer 1

Pydantic V2 已弃用已接受答案中给出的方法。V2 方法是使用自定义序列化器装饰器，因此执行此操作的方法现在如下所示：

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

Answer 2

您并不是试图“加载”任何内容，您的问题是在将模型序列化为JSON 时，您希望使用Enum 名称而不是值对Pydantic 字段进行编码。

您需要做的就是Config向您的BaseModel子类添加一个类，该类指定该Group类型的 JSON 编码器。该json_encoders属性是dict带有序列化器可调用项的键控类型：

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