centralize orjson library use; extend serializer to_json function to return in byte form if specified

Author: shuchenliu

biothings/hub/api/handlers/base.py

@@ -6,10 +6,10 @@
 # see https://github.com/biothings/biothings.api/commit/59c0d78f758018b0d87836657a2b5d1a700503a1
 # import pandas.io.json as pdjson
 # replace pandas json encoder with orjson:
-import orjson
 from tornado.web import RequestHandler
 
 from biothings import config
+from biothings.utils import serializer
 
 
 class DefaultHandler(RequestHandler):
@@ -26,10 +26,9 @@ def write(self, result):
             #     "result": result,
             #     "status": "ok"
             # }, iso_dates=True)
-            orjson.dumps(
+            serializer.to_json({
                 {"result": result, "status": "ok"},
-                option=orjson.OPT_NON_STR_KEYS | orjson.OPT_NAIVE_UTC,
-            ).decode()
+            })
         )
 
     def write_error(self, status_code, **kwargs):

biothings/hub/dataload/dumper.py

@@ -20,6 +20,8 @@
 from typing import Any, Callable, Dict, Generator, Iterable, List, Optional, Tuple, Union
 from urllib import parse as urlparse
 
+from biothings.utils import serializer
+
 try:
     import docker
     from docker.errors import ImageNotFound, NotFound, NullResource
@@ -28,7 +30,6 @@
 except ImportError:
     docker_avail = False
 
-import orjson
 import requests
 
 from biothings import config as btconfig
@@ -1883,7 +1884,7 @@ def _run_api_and_store_to_disk(
     try:
         for filename, obj in fn():
             fn_byte_arr = buffer.setdefault(filename, bytearray())
-            fn_byte_arr.extend(orjson.dumps(obj) + b"\n")
+            fn_byte_arr.extend(serializer.to_json(obj, return_bytes=True) + b"\n")
             if len(fn_byte_arr) >= buffer_size:
                 with open(f"{filename}.{pid}", "ab") as f:
                     f.write(fn_byte_arr)

biothings/utils/dotfield.py

@@ -1,4 +1,3 @@
-import orjson
 from biothings.utils import serializer
 
 
@@ -22,8 +21,7 @@ def make_object(attr, value):
     # s += "}" * (len(attr_list))
     # return json.loads(s)
 
-    # New implementation using orjson module
-    s += orjson.dumps(value).decode("utf-8")  # decoding is necessary because orjson dumps into bytes
+    s += serializer.to_json(value)
     s += "}" * (len(attr_list))
     return serializer.load_json(s)
 

biothings/utils/serializer.py

@@ -28,7 +28,7 @@ def orjson_default(o):
     raise TypeError(f"Type {type(o)} not serializable")
 
 
-def to_json(data, indent=False, sort_keys=False):
+def to_json(data, indent=False, sort_keys=False, return_bytes=False):
     # default option:
     #    OPT_NON_STR_KEYS: non string dictionary key, e.g. integer
     #    OPT_NAIVE_UTC: use UTC as the timezone when it's missing
@@ -37,7 +37,14 @@ def to_json(data, indent=False, sort_keys=False):
         option |= orjson.OPT_INDENT_2
     if sort_keys:
         option |= orjson.OPT_SORT_KEYS
-    return orjson.dumps(data, default=orjson_default, option=option).decode()
+
+    byte_dump = orjson.dumps(data, default=orjson_default, option=option)
+    if return_bytes:
+        return byte_dump
+
+    return byte_dump.decode()
+
+    # return orjson.dumps(data, default=orjson_default, option=option).decode()
 
 
 def to_json_file(data, fobj, indent=False, sort_keys=False):

biothings/web/analytics/channels.py

@@ -2,9 +2,9 @@
 import asyncio
 import certifi
 import logging
-import orjson
 import ssl
 
+from biothings.utils import serializer
 from biothings.web.analytics.events import Event, Message
 
 
@@ -81,7 +81,7 @@ async def send(self, event):
                     "user_id": str(event._cid(1)),
                     "events": events[i : i + 25],
                 }
-                await self.send_request(session, self.url, orjson.dumps(data))
+                await self.send_request(session, self.url, serializer.to_json(data, return_bytes=True))
 
     async def send_request(self, session, url, data):
         retries = 0

tests/web/analytics/test_channels.py

@@ -1,9 +1,10 @@
 import aiohttp
 import asyncio
-import orjson
 import pytest
 
 from aioresponses import aioresponses
+
+from biothings.utils import serializer
 from biothings.web.analytics.channels import SlackChannel, GA4Channel, GAChannel
 from biothings.web.analytics.events import GAEvent, Message
 from unittest.mock import patch
@@ -91,7 +92,8 @@ async def test_send_GA4():
 async def test_send_GA4_request_retries():
     channel = GA4Channel("G-XXXXXX", "SECRET")
     url = channel.url
-    data = orjson.dumps({"test": "data"})
+    # data = orjson.dumps({"test": "data"})
+    data = serializer.to_json({"test": "data"}, return_bytes=True)
 
     async with aiohttp.ClientSession() as session:
         with aioresponses() as responses:
@@ -109,7 +111,7 @@ async def test_send_GA4_request_retries():
 async def test_send_GA4_request_max_retries():
     channel = GA4Channel("G-XXXXXX", "SECRET")
     url = channel.url
-    data = orjson.dumps({"test": "data"})
+    data = serializer.to_json({"test": "data"}, return_bytes=True)
 
     async with aiohttp.ClientSession() as session:
         with aioresponses() as responses: