From bc6aeb46ce6ed8f09c6cdd0054e8a6b4b9981d6d Mon Sep 17 00:00:00 2001
From: Hossein Rashidi <h.rashidi128@gmail.com>
Date: Wed, 7 Aug 2024 15:30:45 +0000
Subject: [PATCH] use orjson for jsonification

---
 aana/api/responses.py |  4 ++--
 aana/storage/op.py    |  6 +++---
 aana/utils/json.py    | 33 +++++++++------------------------
 3 files changed, 14 insertions(+), 29 deletions(-)

diff --git a/aana/api/responses.py b/aana/api/responses.py
index bf589173..e8fb92d8 100644
--- a/aana/api/responses.py
+++ b/aana/api/responses.py
@@ -3,7 +3,7 @@
 import orjson
 from fastapi.responses import JSONResponse
 
-from aana.utils.json import orjson_serializer
+from aana.utils.json import jsonify
 
 
 class AanaJSONResponse(JSONResponse):
@@ -22,4 +22,4 @@ def __init__(self, option: int | None = orjson.OPT_SERIALIZE_NUMPY, **kwargs):
 
     def render(self, content: Any) -> bytes:
         """Override the render method to use orjson.dumps instead of json.dumps."""
-        return orjson_serializer(content, option=self.option)
+        return jsonify(content, option=self.option, as_bytes=True)
diff --git a/aana/storage/op.py b/aana/storage/op.py
index e6f8a03b..fcd413a2 100644
--- a/aana/storage/op.py
+++ b/aana/storage/op.py
@@ -8,7 +8,7 @@
 
 from aana.exceptions.runtime import EmptyMigrationsException
 from aana.utils.core import get_module_dir
-from aana.utils.json import orjson_serializer
+from aana.utils.json import jsonify
 
 
 class DbType(str, Enum):
@@ -30,7 +30,7 @@ def create_postgresql_engine(config):
     connection_string = f"postgresql://{config['user']}:{config['password']}@{config['host']}:{config['port']}/{config['database']}"
     return create_engine(
         connection_string,
-        json_serializer=lambda obj: orjson_serializer(obj).decode(),
+        json_serializer=lambda obj: jsonify(obj),
         json_deserializer=orjson.loads,
     )
 
@@ -47,7 +47,7 @@ def create_sqlite_engine(config):
     connection_string = f"sqlite:///{config['path']}"
     return create_engine(
         connection_string,
-        json_serializer=lambda obj: orjson_serializer(obj).decode(),
+        json_serializer=lambda obj: jsonify(obj),
         json_deserializer=orjson.loads,
     )
 
diff --git a/aana/utils/json.py b/aana/utils/json.py
index e416d976..09308d28 100644
--- a/aana/utils/json.py
+++ b/aana/utils/json.py
@@ -1,13 +1,11 @@
-import json
 from pathlib import Path
 from typing import Any
 
-import numpy as np
 import orjson
 from pydantic import BaseModel
 from sqlalchemy import Engine
 
-__all__ = ["jsonify", "orjson_serializer", "json_serializer_default"]
+__all__ = ["jsonify", "json_serializer_default"]
 
 
 def json_serializer_default(obj: object) -> object:
@@ -38,37 +36,24 @@ def json_serializer_default(obj: object) -> object:
         return str(obj)
     if isinstance(obj, type):
         return str(type)
-    if isinstance(obj, np.ndarray):
-        return str(orjson_serializer(obj))
-    from aana.core.models.media import Media
 
+    from aana.core.models.media import Media
     if isinstance(obj, Media):
         return str(obj)
-    raise TypeError(type(obj))
-
-
-def jsonify(data: Any) -> str:
-    """Convert data to JSON string.
 
-    Args:
-        data (Any): the data
-
-    Returns:
-        str: the JSON string
-    """
-    return json.dumps(data, default=json_serializer_default)
+    raise TypeError(type(obj))
 
 
-def orjson_serializer(
-    content: Any, option: int | None = orjson.OPT_SERIALIZE_NUMPY
-) -> bytes:
+def jsonify(data: Any, option: int | None = orjson.OPT_SERIALIZE_NUMPY, as_bytes: bool = False) -> str | bytes:
     """Serialize content using orjson.
 
     Args:
-        content (Any): The content to serialize.
+        data (Any): The content to serialize.
         option (int | None): The option for orjson.dumps.
+        as_bytes (bool): Return output as bytes instead of string
 
     Returns:
-        bytes: The serialized content.
+        bytes | str: The serialized data as desired format.
     """
-    return orjson.dumps(content, option=option, default=json_serializer_default)
+    output = orjson.dumps(data, option=option, default=json_serializer_default)
+    return output if as_bytes else output.decode()