"""Model definitions for the api module.""" from typing import Any, List, Optional, Dict from flask_sqlalchemy import Pagination from corvus.db import db_model # pylint: disable=too-few-public-methods class APIResponse: """Custom class to wrap api responses.""" def __init__(self, payload: Any, status: int = 200, options: Optional[List[str]] = None) -> None: """Construct an APIResponse object.""" self.payload = payload self.status = status self.options = options # pylint: disable=too-few-public-methods class BaseAPIMessage: """Base class for API responses.""" def to_dict(self) -> Dict[str, Any]: """Abstract to_dict.""" raise NotImplementedError('Not Implemented') # pylint: disable=too-few-public-methods class APIMessage(BaseAPIMessage): """Simple class to encapsulate response messages.""" success: bool message: Optional[str] def __init__(self, success: bool, message: Optional[str]) -> None: """Construct an APIMessage.""" self.success = success self.message = message def to_dict(self) -> Dict[str, Any]: """Serialize an APIMessage to a dict.""" obj: Dict[str, Any] = { 'success': self.success } if self.message is not None: obj['message'] = self.message return obj # pylint: disable=too-few-public-methods class APIPage(BaseAPIMessage): """Simple page response.""" def __init__(self, page: int, total_count: int, last_page: int, items: List[db_model]) -> None: """Construct and APIPage.""" self.page = page self.count = len(items) self.total_count = total_count self.last_page = last_page if last_page > 0 else page self.items = items def to_dict(self) -> Dict[str, Any]: """Serialize an APIPage.""" return { 'page': self.page, 'count': self.count, 'totalCount': self.total_count, 'lastPage': self.last_page, 'items': self.items } def is_empty(self) -> bool: """Check if the page is empty.""" return self.count == 0 and self.total_count == 0 @staticmethod def from_page(page: Pagination) -> 'APIPage': """Create an APIPage from a Pagination object.""" page = APIPage(page.page, page.total, page.pages, page.items) return page if not page.is_empty() else None