Atheneum/server/atheneum/api/decorators.py

"""Decorators to be used in the api module."""
from functools import wraps
from typing import Callable, Any

from flask import jsonify, Response

from atheneum.api.model import APIResponse


def return_json(func: Callable) -> Callable:
    """
    If an Response object is not returned, jsonify the result and return it.

    :param func:
    :return:
    """
    @wraps(func)
    def decorate(*args: list, **kwargs: dict) -> Any:
        """
        Make sure that the return of the function is jsonified into a Response.

        :param args:
        :param kwargs:
        :return:
        """
        result = func(*args, **kwargs)
        if isinstance(result, Response):
            return result
        if isinstance(result, APIResponse):
            return jsonify(result.payload), result.status
        return jsonify(result)

    return decorate