|
|
from enum import Enum |
|
|
from typing import Annotated, Dict, Any, Optional, Type |
|
|
|
|
|
from pydantic import BeforeValidator, BaseModel, ConfigDict, Field |
|
|
from bson import ObjectId |
|
|
|
|
|
PyObjectId = Annotated[str, BeforeValidator(str)] |
|
|
|
|
|
class MongoBaseModel(BaseModel): |
|
|
id: Optional[PyObjectId] = Field(alias="_id", default_factory=ObjectId) |
|
|
|
|
|
model_config = ConfigDict( |
|
|
populate_by_name=True, |
|
|
arbitrary_types_allowed=True, |
|
|
json_encoders={PyObjectId: str}, |
|
|
) |
|
|
|
|
|
def to_mongo(self): |
|
|
result = self.model_dump(by_alias=True, mode='json', exclude={'id',}) |
|
|
result['_id'] = ObjectId(self.id) |
|
|
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 |
|
|
|
