trauma/core/database.py

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