"""JSON specific utilities.""" from datetime import date from typing import Any import rfc3339 from flask.json import JSONEncoder from atheneum.api.model import APIResponse, BaseAPIMessage 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, 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)