Atheneum/server/atheneum/errors.py

"""Error definitions for Atheneum."""
from typing import Dict

from werkzeug.exceptions import HTTPException

from atheneum.api.decorators import return_json
from atheneum.api.model import APIResponse, APIMessage


class BaseError(Exception):
    """Atheneum Base Error Class (5xx errors)."""

    def __init__(
            self,
            message: str = 'Unknown error',
            status_code: int = 500,
            extra_fields: Dict[str, str] = None) -> None:
        """Populate The Error Definition."""
        super().__init__(message)
        self.message = message
        self.status_code = status_code
        self.extra_fields = extra_fields

    def to_dict(self) -> dict:
        """Serialize an error message to return."""
        return {
            'message': self.message,
            'status_code': self.status_code
        }


class ClientError(BaseError):
    """Atheneum errors where the client is wrong (4xx errors)."""

    def __init__(self,
                 message: str = 'Unknown client error',
                 status_code: int = 400,
                 extra_fields: Dict[str, str] = None) -> None:
        """Init for client originated errors."""
        super().__init__(message, status_code, extra_fields)


class ValidationError(ClientError):
    """Atheneum Validation Error."""

    pass


@return_json
def handle_atheneum_404_error(exception: HTTPException) -> APIResponse:
    """Error handler for 404 Atheneum errors."""
    return APIResponse(
        payload=APIMessage(False, 'Not Found'), status=exception.code)


@return_json
def handle_atheneum_base_error(error: BaseError) -> APIResponse:
    """Error handler for basic Atheneum raised errors."""
    return APIResponse(payload=error, status=error.status_code)