Corvus/server/corvus/utility/json_utility.py

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

import rfc3339
from flask.json import JSONEncoder

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