Atheneum/server/atheneum/api/model.py

"""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)