Centralize orjson usage across library (#435)

* replace orjson.loads with serializer method

* remove direct orjson reference from web

* 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,4 @@
-import orjson
+from biothings.utils import serializer
 
 
 def make_object(attr, value):
@@ -21,10 +21,9 @@ 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 orjson.loads(s)
+    return serializer.load_json(s)
 
 
 def merge_object(obj1, obj2):

biothings/utils/parsers.py

@@ -3,7 +3,7 @@
 from typing import Callable, Generator, Iterable, Optional
 from urllib.parse import parse_qsl, urlparse
 
-import orjson
+from biothings.utils import serializer
 
 
 def ndjson_parser(
@@ -31,7 +31,7 @@ def ndjson_parser_func(data_folder):
             for filename in work_dir.glob(pattern):
                 with open(filename, "rb") as f:
                     for line in f:
-                        doc = orjson.loads(line)
+                        doc = serializer.load_json(line)
                         yield doc
 
     return ndjson_parser_func
@@ -60,7 +60,7 @@ def json_array_parser(data_folder):
         for pattern in patterns:
             for filename in work_dir.glob(pattern):
                 with open(filename, "r") as f:
-                    data = orjson.loads(f.read())
+                    data = serializer.load_json(f.read())
                     try:
                         iterator = iter(data)
                     except TypeError:

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,12 @@ 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()
 
 
 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

biothings/web/handlers/base.py

@@ -20,7 +20,6 @@
 """
 import logging
 
-import orjson
 import yaml
 from tornado.web import HTTPError, RequestHandler
 
@@ -105,8 +104,8 @@ def _parse_json(self):
         if not self.request.body:
             return {}
         try:
-            return orjson.loads(self.request.body)
-        except orjson.JSONDecodeError:
+            return serializer.load_json(self.request.body)
+        except serializer.JSONDecodeError:
             raise HTTPError(400, reason="Invalid JSON body.")
 
     def _parse_yaml(self):

biothings/web/options/manager.py

@@ -10,7 +10,7 @@
 from types import MappingProxyType
 
 import jmespath
-import orjson
+from biothings.utils import serializer
 
 try:
     from re import Pattern  # py>=3.7
@@ -250,8 +250,8 @@ def __init__(self, **kwargs):
     def convert_to(self, value, to_type):
         if self.jsoninput:
             try:  # attempt to load as json first
-                _value = orjson.loads(value)
-            except orjson.JSONDecodeError as exc:
+                _value = serializer.load_json(value)
+            except serializer.JSONDecodeError as exc:
                 logging.debug(repr(exc))
             else:  # no more conversions
                 if isinstance(_value, to_type):

biothings/web/query/builder.py

@@ -40,8 +40,8 @@ class implementations or not defined.
 
 from elasticsearch.dsl import MultiSearch, Q, Search
 from elasticsearch.dsl.exceptions import IllegalOperation
-import orjson
 
+from biothings.utils import serializer
 from biothings.utils.common import dotdict
 from biothings.web.query.formatter import ESResultFormatter
 from biothings.web.services.metadata import BiothingsMetadata
@@ -430,9 +430,9 @@ def __init__(self, path):
                             ## alternative implementation  # noqa: E266
                             # self._queries[os.path.basename(dirpath)] = text_file.read()
                             ##
-                            self._queries[os.path.basename(dirpath)] = orjson.loads(text_file.read())
+                            self._queries[os.path.basename(dirpath)] = serializer.load_json(text_file.read())
                         elif "filter" in filename:
-                            self._filters[os.path.basename(dirpath)] = orjson.loads(text_file.read())
+                            self._filters[os.path.basename(dirpath)] = serializer.load_json(text_file.read())
         except Exception:
             self.logger.exception("Error loading user queries.")
 

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: