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.db import db
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
        """Handle encoding date and datetime objects according to rfc3339."""
        try:
            if isinstance(o, date):
                return rfc3339.format(o)
            if isinstance(o, db.Model):
                return serialize_model(o)
            iterable = iter(o)
        except TypeError:
            pass
        else:
            return list(iterable)
        return JSONEncoder.default(self, o)