Atheneum/server/atheneum/utility/json_utility.py

"""JSON specific utilities."""
from datetime import date
from typing import Any

import rfc3339
from flask.json import JSONEncoder

from atheneum.api.model import APIResponse, APIMessage
from atheneum.db import db
from atheneum.errors import BaseError
from atheneum.service.transformation_service import serialize_model


class CustomJSONEncoder(JSONEncoder):
    """Ensure that datetime values are serialized correctly."""

    def default(self, o: Any) -> Any:  # pylint: disable=E0202,R0911
        """Handle encoding date and datetime objects according to rfc3339."""
        try:
            if isinstance(o, date):
                return rfc3339.format(o)
            if isinstance(o, APIResponse):
                payload = o.payload
                if isinstance(payload, db.Model):
                    return serialize_model(o.payload, o.options)
                if isinstance(payload, BaseError):
                    return payload.to_dict()
                if isinstance(payload, APIMessage):
                    return payload.to_dict()
                return payload
            if isinstance(o, db.Model):
                return serialize_model(o)
            iterable = iter(o)
        except TypeError:
            pass
        else:
            return list(iterable)
        return JSONEncoder.default(self, o)