Pydantic v2中__json_schema__失效，如何自定义枚举JSON Schema？

问题背景 #

升级到Pydantic v2后，原v1中通过自定义__json_schema__方法定制Enum类JSON Schema的逻辑失效。原代码通过ExtendedEnum和OneOfEnum的__json_schema__方法，将Enum的Schema从标准enum数组格式，替换为包含oneOf和带title的const项的格式，但v2不再识别该方法，回归了默认的Enum Schema输出。

原v1生成的Schema：

"AssignedSex": {
  "title": "AssignedSex",
  "type": "string",
  "oneOf": [
    {
      "const": "male",
      "title": "Male"
    },
    {
      "const": "female",
      "title": "Female"
    },
    {
      "const": "other",
      "title": "Other"
    },
    {
      "const": "unknown",
      "title": "Unknown"
    }
  ]
}

v2默认生成的Schema：

"AssignedSex": {
  "enum": ["male", "female", "other", "unknown"],
  "title": "AssignedSex",
  "type": "string"
},

解决方案：使用__get_pydantic_json_schema__方法 #

Pydantic v2中，自定义JSON Schema的官方替代方案是实现__get_pydantic_json_schema__类方法，配合__get_pydantic_core_schema__关联验证逻辑与Schema生成规则。针对你的Enum场景，修改后的代码如下：

from enum import Enum
from typing import Optional, List, Dict, Any
from pydantic_core import core_schema
from pydantic import GetJsonSchemaHandler
from pydantic.json_schema import JsonSchemaValue


class ExtendedStr(str):
    _display: str

    def __new__(cls, value: str, display: Optional[str] = None) -> "ExtendedStr":
        s = super().__new__(cls, value)
        s._display = display if display else value
        return s


class ExtendedEnum(ExtendedStr, Enum):
    def __init__(self, value: ExtendedStr) -> None:
        self.value._display = value._display

    @classmethod
    def values(cls) -> List[str]:
        return [c.value for c in cls]

    @classmethod
    def names(cls) -> List[str]:
        return [c.value._display for c in cls]

    @classmethod
    def __get_pydantic_core_schema__(cls, source_type: Any, handler: Any) -> core_schema.CoreSchema:
        # 获取默认核心验证Schema，关联自定义JSON Schema生成方法
        core_schema_obj = handler(source_type)
        return core_schema.no_info_after_validator_function(
            lambda x: x,
            core_schema_obj,
            json_schema=cls.__get_pydantic_json_schema__,
        )

    @classmethod
    def __get_pydantic_json_schema__(cls, schema: JsonSchemaValue, handler: GetJsonSchemaHandler) -> JsonSchemaValue:
        # 先获取默认生成的Schema，再做自定义修改
        schema = handler(schema)
        schema.pop("description", None)
        schema["type"] = "string"
        return schema


class OneOfEnum(ExtendedEnum):
    @classmethod
    def __get_pydantic_json_schema__(cls, schema: JsonSchemaValue, handler: GetJsonSchemaHandler) -> JsonSchemaValue:
        schema = handler(schema)
        # 删除默认enum字段，替换为oneOf结构
        schema.pop("enum", None)
        schema["oneOf"] = [
            {"const": value, "title": name}
            for value, name in zip(cls.values(), cls.names())
        ]
        return schema


class AssignedSex(OneOfEnum):
    male = ExtendedStr("male", "Male")
    female = ExtendedStr("female", "Female")
    other = ExtendedStr("other", "Other")
    unknown = ExtendedStr("unknown", "Unknown")

关键说明 #

• __get_pydantic_core_schema__：定义Pydantic的核心验证逻辑，通过no_info_after_validator_function包装默认Schema，绑定自定义的JSON Schema生成方法，确保验证逻辑正常的同时支持Schema定制。
• __get_pydantic_json_schema__：这是v2替代v1__json_schema__的核心方法，先调用handler(schema)获取默认Schema，再按需修改字段后返回，实现自定义格式。

验证效果 #

修改后的代码生成的Schema将与Pydantic v1完全一致，即包含oneOf和带title的const项结构。

内容的提问来源于stack exchange，提问作者Thomas Erik Olsen