59 lines
1.6 KiB
Python
59 lines
1.6 KiB
Python
from typing import Optional, List
|
|
from pydantic import BaseModel, Field
|
|
from bson import ObjectId
|
|
|
|
class PyObjectId(ObjectId):
|
|
"""
|
|
自定义ObjectId类，用于在Pydantic模型中处理MongoDB的ObjectId
|
|
"""
|
|
@classmethod
|
|
def __get_validators__(cls):
|
|
yield cls.validate
|
|
|
|
@classmethod
|
|
def validate(cls, v, handler):
|
|
if not ObjectId.is_valid(v):
|
|
raise ValueError("Invalid ObjectId")
|
|
return ObjectId(v)
|
|
|
|
@classmethod
|
|
def __get_pydantic_json_schema__(cls, _schema_cache, **_kwargs):
|
|
return {
|
|
'type': 'string',
|
|
'description': 'ObjectId',
|
|
'pattern': r'^[0-9a-fA-F]{24}$'
|
|
}
|
|
|
|
@classmethod
|
|
def __modify_schema__(cls, field_schema):
|
|
field_schema.update(
|
|
type='string',
|
|
description='ObjectId',
|
|
pattern=r'^[0-9a-fA-F]{24}$'
|
|
)
|
|
|
|
class SensorModel(BaseModel):
|
|
"""传感器模型"""
|
|
index: str
|
|
sensor_name: str
|
|
sensor_type: str
|
|
unit: str
|
|
|
|
class DeviceModel(BaseModel):
|
|
id: Optional[PyObjectId] = Field(alias="_id", default=None)
|
|
device_name: str
|
|
device_type: str
|
|
device_number: int
|
|
serial_numbers: list = Field(default_factory=list)
|
|
sensors: List[SensorModel] = Field(default_factory=list)
|
|
|
|
class Config:
|
|
populate_by_name = True
|
|
arbitrary_types_allowed = True
|
|
json_encoders = {ObjectId: str}
|
|
|
|
def dict(self, *args, **kwargs):
|
|
"""确保返回的字典包含 _id 字段"""
|
|
kwargs["by_alias"] = True
|
|
return super().dict(*args, **kwargs)
|
|
|