Pydantic中如何自定义扩展内置集合类型的验证逻辑

在实际开发中，我们经常需要扩展Python内置的集合类型（如tuple、list等）来添加自定义方法。但在使用Pydantic进行数据验证时，这些自定义类型会遇到验证失败的问题。本文将深入探讨如何通过实现__get_pydantic_core_schema__方法来解决这个问题。

问题背景

假设我们有一个文件管理系统，需要处理文件列表数据。我们可能希望创建一个增强版的元组类型，既能保持元组的特性，又能添加自定义方法。例如：

class File(BaseModel):
    name: str
    size: int

class EnhancedFileTuple(tuple[File, ...]):
    def get_total_size(self) -> int:
        return sum(file.size for file in self)

当我们在Pydantic模型中使用这个自定义类型时，会遇到验证错误，因为Pydantic不知道如何处理这个非标准类型。

解决方案

Pydantic提供了__get_pydantic_core_schema__这个类方法，允许我们指定自定义类型的验证逻辑。这个方法需要返回一个核心模式(CoreSchema)，告诉Pydantic如何验证这个类型。

对于扩展内置集合类型的情况，我们可以这样实现：

from pydantic_core import CoreSchema, core_schema
from pydantic import GetCoreSchemaHandler

class EnhancedFileTuple(tuple[File, ...]):
    @classmethod
    def __get_pydantic_core_schema__(
        cls, source_type: Any, handler: GetCoreSchemaHandler
    ) -> CoreSchema:
        # 告诉Pydantic使用标准元组的验证逻辑
        return handler(tuple[File, ...])
    
    def get_total_size(self) -> int:
        return sum(file.size for file in self)

实现原理

1. 核心模式(CoreSchema): 这是Pydantic内部用于描述如何验证和序列化数据的结构。通过实现__get_pydantic_core_schema__，我们可以为自定义类型提供这种描述。

2. 处理程序(handler): 这是一个由Pydantic提供的工具，可以将类型转换为对应的核心模式。通过调用handler并传入基础类型，我们复用Pydantic已有的验证逻辑。

3. 类型保持: 虽然我们使用了标准元组的验证逻辑，但验证后的结果仍会保持为我们的自定义类型，因此自定义方法仍然可用。

实际应用示例

class FileSystemModel(BaseModel):
    recent_files: EnhancedFileTuple

# 可以正常验证
model = FileSystemModel.model_validate({
    "recent_files": [
        {"name": "doc.txt", "size": 1024},
        {"name": "img.jpg", "size": 2048}
    ]
})

# 可以使用自定义方法
print(model.recent_files.get_total_size())  # 输出3072

进阶技巧

1. 添加预处理逻辑：可以在__get_pydantic_core_schema__中定义更复杂的验证逻辑，比如添加数据预处理。

2. 性能考虑：对于高频使用的类型，可以考虑缓存核心模式以避免重复计算。

3. 组合使用：可以结合Pydantic的其他特性，如验证器(validator)和字段类型(Field)，构建更强大的数据模型。

通过这种方式，我们既保持了Pydantic强大的验证能力，又获得了自定义类型的灵活性，是处理复杂数据模型的理想选择。