Merge pull request #62 from solute/add-orjson-handler

Add orjson support

Author: tcalmant

README.md

@@ -64,12 +64,12 @@ A `SimpleJSONRPCServer` class has been added. It is intended to emulate the
 
 ## Requirements
 
-This library supports `ujson`, `cjson` and `simplejson`, and looks for the
-parsers in that order (searching first for `ujson`, `cjson`, `simplejson`
-and finally for the *built-in* `json`).
+This library supports `orjson`, `ujson`, `cjson` and `simplejson`, and looks
+for the parsers in that order (searching first for `orjson`, `ujson`, `cjson`,
+`simplejson` and finally for the *built-in* `json`).
 One of these must be installed to use this library, although if you have a
 standard distribution of 2.7+, you should already have one.
-Keep in mind that `ujson` is supposed to be the quickest, I believe, so if you
+Keep in mind that `orjson` is supposed to be the quickest, I believe, so if you
 are going for full-on optimization you may want to pick it up.
 
 ## Installation

jsonrpclib/jsonlib.py

@@ -103,12 +103,31 @@ def dumps_ujson(obj, encoding="utf-8"):
         return ujson.loads, dumps_ujson
 
 
+class OrJsonHandler(JsonHandler):
+    """
+    Handler based on orjson
+    """
+
+    def get_methods(self):
+        import orjson
+
+        def dumps_orjson(obj, encoding="utf-8"):
+            return orjson.dumps(obj).decode(encoding)
+
+        return orjson.loads, dumps_orjson
+
+
 def get_handler():
     # type: () -> JsonHandler
     """
     Returns the best available Json parser
     """
-    for handler_class in (UJsonHandler, SimpleJsonHandler, CJsonHandler):
+    for handler_class in (
+        OrJsonHandler,
+        UJsonHandler,
+        SimpleJsonHandler,
+        CJsonHandler,
+    ):
         handler = handler_class()
         try:
             loader, dumper = handler.get_methods()

run_tests.sh

@@ -11,6 +11,13 @@ echo "Initial tests..."
 export JSONRPCLIB_TEST_EXPECTED_LIB=json
 coverage run -m pytest || exit 1
 
+echo "orJson tests..."
+pip install orjson && (
+    export JSONRPCLIB_TEST_EXPECTED_LIB=orjson
+    coverage run -m pytest tests/test_jsonlib.py || exit 1
+    pip uninstall -y orjson
+)
+
 echo "uJson tests..."
 pip install ujson && (
     export JSONRPCLIB_TEST_EXPECTED_LIB=ujson

tests/test_jsonlib.py

@@ -109,6 +109,13 @@ def test_best(self):
         except ImportError:
             pass
 
+        try:
+            import orjson  # pylint: disable=unused-import,import-outside-toplevel
+
+            available.add("orjson")
+        except ImportError:
+            pass
+
         # Check the availability of the expected best handler
         expected_best = self._get_expected_best(available)
 
@@ -271,3 +278,47 @@ def test_simplejson(self):
         finally:
             # Reload the module
             imp_reload(simplejson)
+
+    def test_orjson(self):
+        """
+        Tests if the orjson methods are really used
+        """
+        try:
+            import orjson  # pylint: disable=import-outside-toplevel
+        except ImportError:
+            return self.skipTest("orjson is missing: ignore")
+
+        # Check if the expected best is right
+        self._check_expected_best("orjson")
+
+        try:
+            # Force methods to raise an exception
+            orjson.loads = _fake_loads
+
+            orjson.dumps = lambda *args, **kwargs: _fake_dumps(
+                *args, **kwargs
+            ).encode()  # orjson.dumps returns bytes
+
+            # Reload the module
+            imp_reload(jsonlib)
+
+            # Check the handler
+            handler = jsonlib.get_handler()
+            self.assertIsInstance(handler, jsonlib.OrJsonHandler)
+
+            # Check the methods
+            load_method, dump_method = jsonlib.get_handler_methods()
+            self.assertIs(orjson.loads, _fake_loads)
+            self.assertIs(load_method, _fake_loads)
+            self.assertRaises(
+                NotImplementedError, load_method, TEST_INPUT_MARKER
+            )
+            self.assertRaises(
+                NotImplementedError, dump_method, TEST_OUTPUT_MARKER
+            )
+            self.assertRaises(
+                NotImplementedError, dump_method, TEST_OUTPUT_MARKER, "encoding"
+            )
+        finally:
+            # Reload the module
+            imp_reload(orjson)