如何为Pydantic枚举类型添加语义增强？

针对你遇到的Pydantic枚举类及枚举值无法添加语义信息的问题，以下是符合Pydantic 2.x最佳实践的几种解决方案：

方案1：利用枚举类的__json_schema_extra__钩子（官方推荐） #

Pydantic 2.x支持在枚举类中定义__json_schema_extra__类方法，直接为枚举的JSON Schema注入语义描述、自定义字段等信息，无需额外后处理代码。

示例代码 #

from enum import Enum
from pydantic import BaseModel

class OrderStatus(Enum):
    PENDING = "pending"
    SHIPPED = "shipped"
    DELIVERED = "delivered"

    @classmethod
    def __json_schema_extra__(cls, schema: dict, handler) -> None:
        # 为枚举类添加整体语义描述
        schema["description"] = "订单的流转状态，覆盖从创建到完成的全流程节点"
        # 为每个枚举值添加单独的语义说明
        schema["enumDescriptions"] = {
            cls.PENDING.value: "订单已创建，等待商家确认并发货",
            cls.SHIPPED.value: "订单已发出，处于运输配送环节",
            cls.DELIVERED.value: "订单已送达用户，交易完成"
        }
        # 可添加符合语义规范的自定义字段（如Schema.org枚举类型标识）
        schema["@type"] = "Enumeration"

class Order(BaseModel):
    id: str
    status: OrderStatus

# 生成带语义增强的JSON Schema
print(Order.model_json_schema())

效果说明 #

生成的Schema会包含枚举类的description、每个枚举值对应的enumDescriptions，以及自定义的语义字段，完全符合Pydantic的扩展机制，代码耦合度低。

方案2：集中管理语义映射的Schema后处理 #

如果需要统一管理所有枚举的语义信息，避免代码分散在各个枚举类中，可以编写通用的Schema后处理函数，遍历生成的Schema并注入预定义的语义数据。

示例代码 #

from enum import Enum
from pydantic import BaseModel

class OrderStatus(Enum):
    PENDING = "pending"
    SHIPPED = "shipped"
    DELIVERED = "delivered"

class PaymentMethod(Enum):
    ALIPAY = "alipay"
    WECHAT = "wechat"
    CARD = "card"

# 集中存储所有枚举的语义配置
ENUM_SEMANTICS = {
    OrderStatus: {
        "description": "订单流转状态枚举",
        "value_descriptions": {
            OrderStatus.PENDING: "待发货",
            OrderStatus.SHIPPED: "运输中",
            OrderStatus.DELIVERED: "已送达"
        }
    },
    PaymentMethod: {
        "description": "支付方式枚举",
        "value_descriptions": {
            PaymentMethod.ALIPAY: "支付宝支付",
            PaymentMethod.WECHAT: "微信支付",
            PaymentMethod.CARD: "银行卡支付"
        }
    }
}

class Order(BaseModel):
    id: str
    status: OrderStatus
    pay_method: PaymentMethod

def enhance_schema_semantics(schema: dict) -> dict:
    """递归遍历Schema，为枚举类型注入语义信息"""
    def traverse(node):
        if isinstance(node, dict):
            # 匹配枚举类型的Schema节点
            if "enum" in node and "$ref" not in node:
                for enum_cls, semantics in ENUM_SEMANTICS.items():
                    enum_values = [e.value for e in enum_cls]
                    if node["enum"] == enum_values:
                        node["description"] = semantics["description"]
                        node["enumDescriptions"] = {
                            val: semantics["value_descriptions"][enum]
                            for enum, val in zip(enum_cls, enum_values)
                        }
                        break
            # 递归处理子节点
            for key, value in node.items():
                traverse(value)
        elif isinstance(node, list):
            for item in node:
                traverse(item)
        return node

    return traverse(schema.copy())

# 生成并增强Schema
original_schema = Order.model_json_schema()
enhanced_schema = enhance_schema_semantics(original_schema)
print(enhanced_schema)

效果说明 #

这种方式适合枚举数量较多、需要统一维护语义信息的场景，所有语义配置集中在一处，便于修改和管理。

最佳实践总结 #

1. 优先使用官方扩展机制：__json_schema_extra__是Pydantic官方提供的Schema扩展方式，符合框架设计规范，代码更易维护。
2. 集中管理适合多枚举场景：如果项目中有大量枚举需要添加语义，后处理函数能避免代码冗余，提高可维护性。
3. 语义字段尽量遵循规范：自定义语义字段（如enumDescriptions）尽量采用通用命名或遵循Schema.org等语义标准，提升Schema的通用性。

内容的提问来源于stack exchange，提问作者ek365