|
|
from datetime import datetime |
|
|
from enum import Enum |
|
|
from typing import Dict, Any, Type |
|
|
|
|
|
from bson import ObjectId |
|
|
from pydantic import GetCoreSchemaHandler, BaseModel, Field, AnyUrl |
|
|
from pydantic.json_schema import JsonSchemaValue |
|
|
from pydantic_core import core_schema |
|
|
|
|
|
|
|
|
class PyObjectId: |
|
|
@classmethod |
|
|
def __get_pydantic_core_schema__( |
|
|
cls, source: type, handler: GetCoreSchemaHandler |
|
|
) -> core_schema.CoreSchema: |
|
|
return core_schema.with_info_after_validator_function( |
|
|
cls.validate, core_schema.str_schema() |
|
|
) |
|
|
|
|
|
@classmethod |
|
|
def __get_pydantic_json_schema__( |
|
|
cls, schema: core_schema.CoreSchema, handler: GetCoreSchemaHandler |
|
|
) -> JsonSchemaValue: |
|
|
return {"type": "string"} |
|
|
|
|
|
@classmethod |
|
|
def validate(cls, value: str) -> ObjectId: |
|
|
if not ObjectId.is_valid(value): |
|
|
raise ValueError(f"Invalid ObjectId: {value}") |
|
|
return ObjectId(value) |
|
|
|
|
|
def __getattr__(self, item): |
|
|
return getattr(self.__dict__['value'], item) |
|
|
|
|
|
def __init__(self, value: str = None): |
|
|
if value is None: |
|
|
self.value = ObjectId() |
|
|
else: |
|
|
self.value = self.validate(value) |
|
|
|
|
|
def __str__(self): |
|
|
return str(self.value) |
|
|
|
|
|
|
|
|
class MongoBaseModel(BaseModel): |
|
|
id: str = Field(default_factory=lambda: str(PyObjectId())) |
|
|
|
|
|
class Config: |
|
|
arbitrary_types_allowed = True |
|
|
|
|
|
def to_mongo(self) -> Dict[str, Any]: |
|
|
def model_to_dict(model: BaseModel) -> Dict[str, Any]: |
|
|
doc = {} |
|
|
for name, value in model._iter(): |
|
|
key = model.__fields__[name].alias or name |
|
|
|
|
|
if isinstance(value, BaseModel): |
|
|
doc[key] = model_to_dict(value) |
|
|
elif isinstance(value, list) and all(isinstance(i, BaseModel) for i in value): |
|
|
doc[key] = [model_to_dict(item) for item in value] |
|
|
elif value and isinstance(value, Enum): |
|
|
doc[key] = value.value |
|
|
elif isinstance(value, datetime): |
|
|
doc[key] = value.isoformat() |
|
|
elif value and isinstance(value, AnyUrl): |
|
|
doc[key] = str(value) |
|
|
else: |
|
|
doc[key] = value |
|
|
|
|
|
return doc |
|
|
|
|
|
result = model_to_dict(self) |
|
|
return result |
|
|
|
|
|
@classmethod |
|
|
def from_mongo(cls, data: Dict[str, Any]): |
|
|
def restore_enums(inst: Any, model_cls: Type[BaseModel]) -> None: |
|
|
for name, field in model_cls.__fields__.items(): |
|
|
value = getattr(inst, name) |
|
|
if field and isinstance(field.annotation, type) and issubclass(field.annotation, Enum): |
|
|
setattr(inst, name, field.annotation(value)) |
|
|
elif isinstance(value, BaseModel): |
|
|
restore_enums(value, value.__class__) |
|
|
elif isinstance(value, list): |
|
|
for i, item in enumerate(value): |
|
|
if isinstance(item, BaseModel): |
|
|
restore_enums(item, item.__class__) |
|
|
elif isinstance(field.annotation, type) and issubclass(field.annotation, Enum): |
|
|
value[i] = field.annotation(item) |
|
|
elif isinstance(value, dict): |
|
|
for k, v in value.items(): |
|
|
if isinstance(v, BaseModel): |
|
|
restore_enums(v, v.__class__) |
|
|
elif isinstance(field.annotation, type) and issubclass(field.annotation, Enum): |
|
|
value[k] = field.annotation(v) |
|
|
|
|
|
if data is None: |
|
|
return None |
|
|
instance = cls(**data) |
|
|
restore_enums(instance, instance.__class__) |
|
|
return instance |
|
|
|
