Refactor/pick a faster json lib (#1152)

* use rapidjson instead

* use orjson instead

* fix tests; relocate some report properties

* fix pyproject.toml

* add newsfrag; some tweaks

doc/newsfragments/3147_changed.another_json_lib.rst

@@ -0,0 +1 @@
+Use a new JSON library ``orjson`` to improve performance when using Python 3.8 or later versions.

examples/ExecutionPools/Discover/test_plan.py

@@ -45,6 +45,7 @@ def main(plan):
 
 if __name__ == "__main__":
     res = main()
-    assert len(res.report.entries) == 5
+    if res.report.entries:
+        assert len(res.report.entries) == 5
     print("Exiting code: {}".format(res.exit_code))
     sys.exit(res.exit_code)

pyproject.toml

@@ -66,7 +66,9 @@
         "typing_extensions",
         "dill",
         "gherkin-official==4.1.3",
-        "parse"
+        "parse",
+        "orjson; python_version>='3.8'",
+        "flask-orjson; python_version>='3.8'"
     ]
     requires-python = ">=3.7"
 

testplan/common/report/base.py

@@ -263,8 +263,8 @@ def __init__(
         uid: Optional[str] = None,
         entries: Optional[list] = None,
         parent_uids: Optional[List[str]] = None,
-        status_override=None,
-        status_reason=None,
+        status_override: Optional[Status] = None,
+        status_reason: Optional[str] = None,
     ):
         self.name = name
         self.description = description
@@ -433,6 +433,56 @@ def is_empty(self) -> bool:
         """
         return len(self.entries) == len(self.logs) == 0
 
+    @property
+    def passed(self) -> bool:
+        """Shortcut for getting if report status should be considered passed."""
+        return self.status.normalised() == Status.PASSED
+
+    @property
+    def failed(self) -> bool:
+        """
+        Shortcut for checking if report status should be considered failed.
+        """
+        return self.status <= Status.FAILED
+
+    @property
+    def unstable(self) -> bool:
+        """
+        Shortcut for checking if report status should be considered unstable.
+        """
+        return self.status.normalised() == Status.UNSTABLE
+
+    @property
+    def unknown(self) -> bool:
+        """
+        Shortcut for checking if report status is unknown.
+        """
+        return self.status.normalised() == Status.UNKNOWN
+
+    @property
+    def status(self) -> Status:
+        """Return the report status."""
+        if self.status_override:
+            return self.status_override
+        return self._status
+
+    @status.setter
+    def status(self, new_status: Status):
+        self._status = new_status
+
+    @property
+    def runtime_status(self) -> RuntimeStatus:
+        """
+        Used for interactive mode, the runtime status of a testcase will be one
+        of ``RuntimeStatus``.
+        """
+        return self._runtime_status
+
+    @runtime_status.setter
+    def runtime_status(self, new_status: RuntimeStatus):
+        """Set the runtime status."""
+        self._runtime_status = new_status
+
     @property
     def hash(self):
         """Return a hash of all entries in this report."""
@@ -468,34 +518,8 @@ def __init__(self, name, **kwargs):
         for child in self.entries:
             self.set_parent_uids(child)
 
-    @property
-    def passed(self):
-        """Shortcut for getting if report status should be considered passed."""
-        return self.status.normalised() == Status.PASSED
-
-    @property
-    def failed(self):
-        """
-        Shortcut for checking if report status should be considered failed.
-        """
-        return self.status <= Status.FAILED
-
-    @property
-    def unstable(self):
-        """
-        Shortcut for checking if report status should be considered unstable.
-        """
-        return self.status.normalised() == Status.UNSTABLE
-
-    @property
-    def unknown(self):
-        """
-        Shortcut for checking if report status is unknown.
-        """
-        return self.status.normalised() == Status.UNKNOWN
-
-    @property
-    def status(self):
+    @Report.status.getter
+    def status(self) -> Status:
         """
         Status of the report, will be used to decide
         if a Testplan run has completed successfully or not.
@@ -513,12 +537,8 @@ def status(self):
 
         return self._status
 
-    @status.setter
-    def status(self, new_status):
-        self._status = new_status
-
     @property
-    def runtime_status(self):
+    def runtime_status(self) -> RuntimeStatus:
         """
         The runtime status is used for interactive running, and reports
         whether a particular entry is READY, WAITING, RUNNING, RESETTING,
@@ -534,7 +554,7 @@ def runtime_status(self):
         return self._runtime_status
 
     @runtime_status.setter
-    def runtime_status(self, new_status):
+    def runtime_status(self, new_status: RuntimeStatus):
         """Set the runtime_status of all child entries."""
         for entry in self:
             if entry.category != ReportCategories.SYNTHESIZED:
@@ -635,11 +655,11 @@ def remove_by_uid(self, uid):
 
     __delitem__ = remove_by_uid
 
-    def pre_order_reports(self):
+    def pre_order_iterate(self):
         yield self
         for e in self:
             if isinstance(e, BaseReportGroup):
-                yield from e.pre_order_reports()
+                yield from e.pre_order_iterate()
             elif isinstance(e, Report):
                 yield e
 
@@ -961,7 +981,7 @@ def set_runtime_status_filtered(
     ) -> None:
         """
         Alternative setter for the runtime status of an entry. Propagates only
-          to the specified entries.
+        to the specified entries.
 
         :param new_status: new runtime status to be set
         :param entries: tree-like structure of entries names

testplan/common/report/schemas.py

@@ -4,12 +4,16 @@
 from marshmallow import Schema, fields, post_load
 from marshmallow.utils import EXCLUDE
 
+from testplan.common.report.base import (
+    BaseReportGroup,
+    Report,
+    RuntimeStatus,
+    Status,
+)
 from testplan.common.serialization import fields as custom_fields
 from testplan.common.serialization import schemas
 from testplan.common.utils import timing
 
-from .base import Report, BaseReportGroup, Status, RuntimeStatus
-
 __all__ = ["ReportLogSchema", "ReportSchema", "BaseReportGroupSchema"]
 
 # pylint: disable=unused-argument
@@ -91,6 +95,14 @@ class Meta:
         allow_none=True,
     )
     status_reason = fields.String(allow_none=True)
+    status = fields.Function(
+        lambda x: x.status.to_json_compatible(),
+        Status.from_json_compatible,
+    )
+    runtime_status = fields.Function(
+        lambda x: x.runtime_status.to_json_compatible(),
+        RuntimeStatus.from_json_compatible,
+    )
     logs = fields.Nested(ReportLogSchema, many=True)
     hash = fields.Integer(dump_only=True)
     parent_uids = fields.List(fields.String())
@@ -127,14 +139,6 @@ class BaseReportGroupSchema(ReportSchema):
         },
         many=True,
     )
-    status = fields.Function(
-        lambda x: x.status.to_json_compatible(),
-        Status.from_json_compatible,
-    )
-    runtime_status = fields.Function(
-        lambda x: x.runtime_status.to_json_compatible(),
-        RuntimeStatus.from_json_compatible,
-    )
     counter = fields.Dict(dump_only=True)
     children = fields.List(fields.Nested(ReportLinkSchema))
 

testplan/common/utils/json.py

@@ -0,0 +1,43 @@
+import json
+from pathlib import Path
+from typing import Union
+
+_USE_ORJSON = False
+
+try:
+    import orjson
+except ImportError:
+    pass
+else:
+    _USE_ORJSON = True
+
+
+def json_loads(data: str):
+    if _USE_ORJSON:
+        return orjson.loads(data)
+    else:
+        return json.loads(data)
+
+
+def json_dumps(data, indent_2=False, default=None) -> str:
+    if _USE_ORJSON:
+        return orjson.dumps(
+            data,
+            default=default,
+            option=orjson.OPT_INDENT_2 if indent_2 else 0,
+        ).decode()
+    else:
+        if default:
+
+            class _E(json.JSONEncoder):
+                def default(self, o):
+                    return default(o)
+
+        else:
+            _E = None
+        return json.dumps(data, cls=_E, indent=2 if indent_2 else None)
+
+
+def json_load_from_path(path: Union[str, Path]) -> dict:
+    with open(path) as fp:
+        return json_loads(fp.read())

testplan/exporters/testing/http/__init__.py

@@ -3,7 +3,6 @@
 must be able to handle POST request and receive data in JSON format.
 """
 
-import json
 from typing import Any, Tuple, Union, Optional, Dict
 
 import requests
@@ -15,19 +14,13 @@
     ExportContext,
     verify_export_context,
 )
+from testplan.common.utils.json import json_dumps
 from testplan.common.utils.validation import is_valid_url
 from testplan.report import TestReport
 from testplan.report.testing.schemas import TestReportSchema
 from ..base import Exporter
 
 
-class CustomJsonEncoder(json.JSONEncoder):
-    """To jsonify data that cannot be serialized by default JSONEncoder."""
-
-    def default(self, obj: Any) -> str:  # pylint: disable = method-hidden
-        return str(obj)
-
-
 class HTTPExporterConfig(ExporterConfig):
     """
     Configuration object for
@@ -83,7 +76,7 @@ def _upload_report(
                 response = requests.post(
                     url=url,
                     headers=headers,
-                    data=json.dumps(data, cls=CustomJsonEncoder),
+                    data=json_dumps(data, default=str),
                     timeout=self.cfg.timeout,
                 )
                 response.raise_for_status()

testplan/exporters/testing/json/base.py

@@ -4,7 +4,6 @@
 """
 
 import hashlib
-import json
 import os
 import pathlib
 
@@ -17,6 +16,7 @@
     ExportContext,
     verify_export_context,
 )
+from testplan.common.utils.json import json_dumps, json_loads
 from testplan.common.utils.path import makedirs
 from testplan.defaults import ATTACHMENTS, RESOURCE_DATA
 from testplan.report.testing.base import TestReport, TestCaseReport
@@ -60,7 +60,7 @@ def save_resource_data(
 ) -> pathlib.Path:
     directory.mkdir(parents=True, exist_ok=True)
     with open(report.resource_meta_path) as meta_file:
-        meta_info = json.load(meta_file)
+        meta_info = json_loads(meta_file.read())
     for host_meta in meta_info["entries"]:
         if "resource_file" in host_meta:
             dist_path = (
@@ -70,7 +70,7 @@ def save_resource_data(
             host_meta["resource_file"] = dist_path.name
     meta_path = directory / pathlib.Path(report.resource_meta_path).name
     with open(meta_path, "w") as meta_file:
-        json.dump(meta_info, meta_file)
+        meta_file.write(json_dumps(meta_info))
     return meta_path
 
 
@@ -172,9 +172,9 @@ def export(
                 attachments_dir.mkdir(parents=True, exist_ok=True)
 
                 with open(structure_filepath, "w") as json_file:
-                    json.dump(structure, json_file)
+                    json_file.write(json_dumps(structure))
                 with open(assertions_filepath, "w") as json_file:
-                    json.dump(assertions, json_file)
+                    json_file.write(json_dumps(assertions))
 
                 meta["attachments"] = save_attachments(
                     report=source, directory=attachments_dir
@@ -190,15 +190,15 @@ def export(
                 meta["assertions_file"] = assertions_filename
 
                 with open(json_path, "w") as json_file:
-                    json.dump(meta, json_file)
+                    json_file.write(json_dumps(meta))
             else:
                 data["attachments"] = save_attachments(
                     report=source, directory=attachments_dir
                 )
                 data["version"] = 1
 
                 with open(json_path, "w") as json_file:
-                    json.dump(data, json_file)
+                    json_file.write(json_dumps(data))
 
             self.logger.user_info("JSON generated at %s", json_path)
             result = {"json": self.cfg.json_path}

testplan/importers/testplan.py

@@ -1,11 +1,11 @@
 """
 Implements one-phase importer for Testplan JSON format.
 """
-import json
 from typing import List
 
-from testplan.importers import ResultImporter, ImportedResult
-from testplan.report import TestGroupReport, TestReport, ReportCategories
+from testplan.common.utils.json import json_loads
+from testplan.importers import ImportedResult, ResultImporter
+from testplan.report import ReportCategories, TestGroupReport, TestReport
 from testplan.report.testing.schemas import TestReportSchema
 
 
@@ -41,7 +41,7 @@ def __init__(self, path: str):
     def import_result(self) -> ImportedResult:
         """ """
         with open(self.path) as fp:
-            result_json = json.load(fp)
+            result_json = json_loads(fp.read())
             result = self.schema.load(result_json)
 
             return TestplanImportedResult(result)