"""Model definitions for the api module.""" from typing import Any, List, Optional, Dict, Type from flask_sqlalchemy import Pagination from atheneum import db # 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[Type[db.Model]]) -> None: """Construct and APIPage.""" self.page = page self.count = len(items) self.total_count = total_count self.last_page = last_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 } @staticmethod def from_page(page: Pagination) -> 'APIPage': """Create an APIPage from a Pagination object.""" return APIPage(page.page, page.total, page.pages, page.items)