| 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 |
| |
