test: refactor benchmark test to support more json libraries

test/benchmark_test.py

@@ -4,6 +4,21 @@
 
 
 class TestBenchmark(unittest.TestCase):
+    def setUp(self):
+        try:
+            import orjson
+            self._orjson = orjson
+        except ImportError:
+            print("orjson is not installed")
+            self._orjson = None
+
+        try:
+            import ujson
+            self._ujson = ujson
+        except ImportError:
+            print("ujson is not installed")
+            self._ujson = None
+
     def time_benchmark(self, repeat_time, func, *args, **kwargs):
         import time
         time_0 = time.time()
@@ -12,6 +27,12 @@ def time_benchmark(self, repeat_time, func, *args, **kwargs):
         time_1 = time.time()
         return time_1 - time_0
 
+    def _run_test_setup(self, calling_setups, data, base_name, test_name):
+        for setup in calling_setups:
+            name, func, kwargs = setup
+            time_ = self.time_benchmark(1000, func, data, **kwargs)
+            print(f"test {test_name}, file: {base_name}, time_{name}: {time_}")
+
     def test_benchmark_encode(self):
         import json
 
@@ -21,14 +42,25 @@ def test_benchmark_encode(self):
 
         bench_files = get_benchfiles_fullpath()
 
+        std_calling_kwargs = {
+            "ensure_ascii": False
+        }
+        std_json_setup = ('std', json.dumps, std_calling_kwargs)
+        cjson_setup = ('cjson', cjson.dumps, {})
+
+        calling_setups = [std_json_setup, cjson_setup]
+        if self._orjson is not None:
+            orjson_setup = ('orjson', self._orjson.dumps, {})
+            calling_setups.append(orjson_setup)
+        if self._ujson is not None:
+            ujson_setup = ('ujson', self._ujson.dumps, {})
+            calling_setups.append(ujson_setup)
+
         for filename in bench_files:
             base_name = os.path.basename(filename)
             with open(filename, "r", encoding='utf-8') as f:
                 data = json.load(f)
-            time_std = self.time_benchmark(1000, json.dumps, data, indent=None, ensure_ascii=False)
-            print(f"test encode, file: {base_name}, time_std: {time_std}")
-            time_cjson = self.time_benchmark(1000, cjson.dumps, data)
-            print(f"test encode, file: {base_name}, time_cjson: {time_cjson}")
+            self._run_test_setup(calling_setups, data, base_name, 'encode')
 
     def test_benchmark_decode(self):
         import json
@@ -39,14 +71,21 @@ def test_benchmark_decode(self):
 
         bench_files = get_benchfiles_fullpath()
 
+        std_json_setup = ('std', json.loads, {})
+        cjson_setup = ('cjson', cjson.loads, {})
+        calling_setups = [std_json_setup, cjson_setup]
+        if self._orjson is not None:
+            orjson_setup = ('orjson', self._orjson.loads, {})
+            calling_setups.append(orjson_setup)
+        if self._ujson is not None:
+            ujson_setup = ('ujson', self._ujson.loads, {})
+            calling_setups.append(ujson_setup)
+
         for filename in bench_files:
             base_name = os.path.basename(filename)
             with open(filename, "r", encoding='utf-8') as f:
                 data = f.read()
-            time_std = self.time_benchmark(1000, json.loads, data)
-            print(f"test decode, file: {base_name}, time_std: {time_std}")
-            time_cjson = self.time_benchmark(1000, cjson.loads, data)
-            print(f"test decode, file: {base_name}, time_cjson: {time_cjson}")
+            self._run_test_setup(calling_setups, data, base_name, 'decode')
 
 
 if __name__ == "__main__":