Spaces:
Sleeping
Sleeping
| from datetime import date, datetime, time | |
| from enum import Enum | |
| from typing import Any, Dict, Type | |
| from bson import ObjectId | |
| from pydantic import AnyUrl, BaseModel, Field, GetCoreSchemaHandler | |
| from pydantic.json_schema import JsonSchemaValue | |
| from pydantic_core import core_schema | |
| class PyObjectId: | |
| 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() # type: ignore | |
| ) | |
| def __get_pydantic_json_schema__( | |
| cls, _schema: core_schema.CoreSchema, _handler: GetCoreSchemaHandler | |
| ) -> JsonSchemaValue: | |
| return {"type": "string"} | |
| 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 = None): | |
| if value is None: | |
| self.value = ObjectId() | |
| else: | |
| self.value = self.validate(value) | |
| def __str__(self): | |
| return str(self.value) | |
| class MongoSerializableModel(BaseModel): | |
| def to_mongo(self) -> Dict[str, Any]: | |
| def model_to_dict(model: BaseModel) -> Dict[str, Any]: | |
| doc = {} | |
| for name in model.__fields__.keys(): | |
| value = getattr(model, name) | |
| 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] # type: ignore | |
| elif value and isinstance(value, Enum): | |
| doc[key] = value.value | |
| elif isinstance(value, (datetime, date, time)): | |
| doc[key] = value.isoformat() # type: ignore | |
| elif value and isinstance(value, AnyUrl): | |
| doc[key] = str(value) # type: ignore | |
| else: | |
| doc[key] = value | |
| return doc | |
| return model_to_dict(self) | |
| class MongoBaseModel(MongoSerializableModel): | |
| id: str = Field(default_factory=lambda: str(PyObjectId())) | |
| class Config: | |
| arbitrary_types_allowed = True | |
| extra = "ignore" | |
| populate_by_name = True | |
| validate_assignment = True | |
| 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(): # type: ignore | |
| 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 | |