30 lines
878 B
Python
30 lines
878 B
Python
from bson import ObjectId
|
|
from pydantic import GetCoreSchemaHandler
|
|
from pydantic_core import core_schema
|
|
|
|
|
|
class PyObjectId(ObjectId):
|
|
"""Custom ObjectId type for Pydantic v2 to handle MongoDB _id"""
|
|
|
|
@classmethod
|
|
def __get_pydantic_core_schema__(
|
|
cls, source, handler: GetCoreSchemaHandler
|
|
) -> core_schema.CoreSchema:
|
|
return core_schema.no_info_after_validator_function(
|
|
cls.validate,
|
|
core_schema.union_schema(
|
|
[
|
|
core_schema.str_schema(),
|
|
core_schema.is_instance_schema(ObjectId),
|
|
]
|
|
),
|
|
)
|
|
|
|
@classmethod
|
|
def validate(cls, v):
|
|
if isinstance(v, ObjectId):
|
|
return v
|
|
if isinstance(v, str) and ObjectId.is_valid(v):
|
|
return ObjectId(v)
|
|
raise ValueError(f"Invalid ObjectId: {v}")
|