kraken/backend/schemas/trains.py

"""Module schemas/trains verifies table data type"""
from datetime import datetime
from typing import Literal, Optional
from typing import Union

from pydantic import BaseModel, Field

from schemas.hand_peripheral import ControlMotion


def get_timestamp() -> datetime:
    return datetime.strptime(datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                             "%Y-%m-%d %H:%M:%S")


class TrainBase(BaseModel):
    position: Optional[str] = None
    rank: Optional[str] = None
    trial_num: Optional[str] = None


class TrainUpdate(TrainBase):
    # position: str
    # rank: str
    # trial_num: int
    start_time: Optional[datetime] = Field(default_factory=get_timestamp)
    end_time: Optional[datetime] = Field(default_factory=get_timestamp)


class TrainCreate(TrainBase):
    # position: str
    # rank: str
    # trial_num: int
    start_time: Optional[datetime] = Field(default_factory=get_timestamp)
    end_time: Optional[datetime] = Field(default_factory=get_timestamp)
    device_param: Union[ControlMotion, None] = None
    owner_id: str


class ShowTrain(TrainBase):
    position: str
    rank: str
    trial_num: int
    start_time: datetime
    end_time: datetime
    grade: str = None
    consume_time: int = None
    accuracy: float = None
    is_train: bool = False
    medical_certificate: str = ""

    class Config():
        orm_mode = True

class ShowTrainWithVideo(ShowTrain):
    video_path: list


class TrainResult(BaseModel):
    grade: Literal["优秀", "良好", "尚可"] = Field()
    accuracy: float = Field()
    consume_time: float = Field()


class TrainMedicalCertificate(BaseModel):
    medical_certificate: str