Atheneum/server/atheneum/api/model.py

"""Model definitions for the api module."""
from typing import Any, List, Optional, Dict


class APIResponse:  # pylint: disable=too-few-public-methods
    """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


class APIMessage:  # pylint: disable=too-few-public-methods
    """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:
        """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