schemas/dataset.py

from datetime import datetime

from pydantic import BaseModel

from schemas.document import Document


class Dataset(BaseModel):
    """
    Краткая информация о датасете.
    """

    id: int
    dateCreated: datetime
    name: str
    isActive: bool
    isDraft: bool

    class Config:
        from_attributes = True
        json_schema_extra = {
            'example': {
                'id': 1,
                'dateCreated': datetime.fromisoformat('2024-01-01T00:00:00'),
                'name': 'default',
                'isActive': True,
                'isDraft': False,
            }
        }


class DocumentsPage(BaseModel):
    """
    Страница с документами.
    """

    page: list[Document]
    total: int
    pageNumber: int
    pageSize: int

    class Config:
        json_schema_extra = {
            'example': {
                'documents': [Document.Config.json_schema_extra['example']],
                'total': 500,
                'pageNumber': 1,
                'pageSize': 10,
            }
        }


class DatasetExpanded(Dataset):
    """
    Расширенная информация о датасете, включающая страницу с документами.
    """

    data: DocumentsPage

    class Config:
        json_schema_extra = {
            'example': {
                'id': 1,
                'dateCreated': datetime.fromisoformat('2024-01-01T00:00:00'),
                'name': 'default',
                'isActive': True,
                'isDraft': False,
                'data': DocumentsPage.Config.json_schema_extra['example'],
            }
        }


class DatasetProcessing(BaseModel):
    """
    Информация о процессе обработки датасета.
    """

    status: str
    total: int | None
    current: int | None
    datasetName: str | None

    class Config:
        json_schema_extra = {
            'example': {
                'status': 'ready',
            }
        }


class SortQuery(BaseModel):
    """
    Запрос на сортировку.
    """

    field: str
    direction: str

    class Config:
        json_schema_extra = {'example': {'field': 'name', 'direction': 'asc'}}


class SortQueryList(BaseModel):
    """
    Список запросов на сортировку.
    """

    sorts: list[SortQuery]

    class Config:
        json_schema_extra = {'example': [{'field': 'name', 'direction': 'asc'}]}