在 Pydantic v2 中使用 bson.ObjectId

Question

我找到了一些示例，了解如何在

BaseModel

类中使用ObjectId。基本上，这可以通过创建 Pydantic 友好的类来实现，如下所示：

class PyObjectId(ObjectId):
    @classmethod
    def __get_validators__(cls):
        yield cls.validate

    @classmethod
    def validate(cls, v):
        if not ObjectId.is_valid(v):
            raise ValueError("Invalid objectid")
        return ObjectId(v)

    @classmethod
    def __modify_schema__(cls, field_schema):
        field_schema.update(type="string")

但是，这似乎适用于 Pydantic v1，因为该机制已被

__get_pydantic_core_schema__

类方法取代。但是，我无法使用 Pydantic v2 实现等效的解决方案。有可能吗？我需要什么验证器？我尝试重构一些东西，但无法得到任何可用的东西。

Answer 1

要将老式的

PyObjectId

迁移到最新的 pydantic-v2 版本，最简单的方法是使用 带注释的验证器。

from typing import Any
from typing import Annotated, Union
from bson import ObjectId
from pydantic import PlainSerializer, AfterValidator, WithJsonSchema

def validate_object_id(v: Any) -> ObjectId:
    if isinstance(v, ObjectId):
        return v
    if ObjectId.is_valid(v):
        return ObjectId(v)
    raise ValueError("Invalid ObjectId")

PyObjectId = Annotated[
    Union[str, ObjectId],
    AfterValidator(validate_object_id),
    PlainSerializer(lambda x: str(x), return_type=str),
    WithJsonSchema({"type": "string"}, mode="serialization"),
]

然后您可以通过以下方式在模型中使用它：

from pydantic import BaseModel
from pydantic import ConfigDict, Field

class MyCustomModel(BaseModel):
    id: PyObjectId = Field(alias="_id")

    model_config = ConfigDict(arbitrary_types_allowed=True)

使用 TypeAdapter 进行测试：

import pytest
from bson import ObjectId
from pydantic import TypeAdapter, ConfigDict
    
@pytest.mark.parametrize("obj", ["64b7992ba8f08069073f1055", ObjectId("64b7992ba8f08069073f1055")])
def test_pyobjectid_validation(obj):
    ta = TypeAdapter(PyObjectId, config=ConfigDict(arbitrary_types_allowed=True))
    ta.validate_python(obj)

@pytest.mark.parametrize("obj", ["64b7992ba8f08069073f1055", ObjectId("64b7992ba8f08069073f1055")])
def test_pyobjectid_serialization(obj):
    ta = TypeAdapter(PyObjectId, config=ConfigDict(arbitrary_types_allowed=True))
    ta.dump_json(obj)

即使使用最新的 FastAPI v0.100.0+

，该解决方案也能正常工作

Answer 2

解决此 Pydantic-V2 问题的最简单方法是创建一个 带注释的验证器：

from typing_extensions import Annotated

from pydantic import BaseModel, ValidationError, field_validator
from pydantic.functional_validators import AfterValidator


PyObjectId = Annotated[
    ObjectId,
    AfterValidator(ObjectId.is_valid),
]

就表示为字符串而言，您可能想要更新注释，可能使用带有字符串的 Union，然后执行转换为字符串的附加验证器：

PyObjectId = Annotated[
    Union[ObjectId, str],
    AfterValidator(ObjectId.is_valid),
    AfterValidator(str)
]

使用 TypeAdaptor 进行测试：

from pydantic import TypeAdapter


ta = TypeAdapter(PyObjectId)
ta.validate_python("SOME-STRING")

Answer 3

以上都不适合我。我按照 Pydantic 文档 提出了这个解决方案：

from typing import Annotated, Any, Callable

from bson import ObjectId
from fastapi import FastAPI
from pydantic import BaseModel, ConfigDict, Field, GetJsonSchemaHandler
from pydantic.json_schema import JsonSchemaValue
from pydantic_core import core_schema


# Based on https://docs.pydantic.dev/latest/usage/types/custom/#handling-third-party-types
class _ObjectIdPydanticAnnotation:
    @classmethod
    def __get_pydantic_core_schema__(
            cls,
            _source_type: Any,
            _handler: Callable[[Any], core_schema.CoreSchema],
    ) -> core_schema.CoreSchema:

        def validate_from_str(id_: str) -> ObjectId:
            return ObjectId(id_)

        from_str_schema = core_schema.chain_schema(
            [
                core_schema.str_schema(),
                core_schema.no_info_plain_validator_function(validate_from_str),
            ]
        )

        return core_schema.json_or_python_schema(
            json_schema=from_str_schema,
            python_schema=core_schema.union_schema(
                [
                    # check if it's an instance first before doing any further work
                    core_schema.is_instance_schema(ObjectId),
                    from_str_schema,
                ]
            ),
            serialization=core_schema.plain_serializer_function_ser_schema(
                lambda instance: str(instance)
            ),
        )

    @classmethod
    def __get_pydantic_json_schema__(
            cls, _core_schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler
    ) -> JsonSchemaValue:
        # Use the same schema that would be used for `str`
        return handler(core_schema.str_schema())


PydanticObjectId = Annotated[
    ObjectId, _ObjectIdPydanticAnnotation
]


class User(BaseModel):
    model_config = ConfigDict(populate_by_name=True)

    id: PydanticObjectId = Field(alias='_id')
    name: str


app = FastAPI()


@app.get("/user/{id}")
def get_usr(id: str) -> User:
    # Here we would connect to MongoDB and return the user.
    # Method is here just to test that FastAPI will not complain about the "User" return type
    pass


# Some usage examples

user1 = User(_id=ObjectId('64cca8a68efc81fc425aa864'), name='John Doe')
user2 = User(_id='64cca8a68efc81fc425aa864', name='John Doe')
assert user1 == user2  # Can use str and ObjectId interchangeably

# Serialization
assert repr(user1) == "User(id=ObjectId('64cca8a68efc81fc425aa864'), name='John Doe')"
assert user1.model_dump() == {'id': '64cca8a68efc81fc425aa864', 'name': 'John Doe'}
assert user1.model_dump_json() == '{"id":"64cca8a68efc81fc425aa864","name":"John Doe"}'

# Deserialization
user2 = User.model_validate_json('{"id":"64cca8a68efc81fc425aa864","name":"John Doe"}')
user3 = User.model_validate_json('{"_id":"64cca8a68efc81fc425aa864","name":"John Doe"}')
assert user1 == user2 == user3

user4 = User(_id=ObjectId(), name='Jane Doe')  # Default ObjectId constructor

# Validation
user5 = User(_id=ObjectId('qwe'), name='Jack Failure')  # Will throw bson.errors.InvalidId