Corvus/server/server/utility/json_utility.py

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

import rfc3339
from flask.json import JSONEncoder

from server.api.model import APIResponse, BaseAPIMessage
from server.db import db
from server.errors import BaseError
from server.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, BaseAPIMessage):
                    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)