File size: 2,086 Bytes
4cd2145 de09008 4cd2145 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
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
|