Improve json performance by porting core orjson utils (#4816)

* Improve json performance by porting core orjson utils

* port relevant tests

* pylint

* add test for read_json_file

* add test for read_json_file

* remove workaround for core issue we do not have here

---------

Co-authored-by: Pascal Vizeli <pvizeli@syshack.ch>

requirements.txt

@@ -17,6 +17,7 @@ docker==7.0.0
 faust-cchardet==2.1.19
 gitpython==3.1.41
 jinja2==3.1.3
+orjson==3.9.10
 pulsectl==23.5.2
 pyudev==0.24.1
 PyYAML==6.0.1

supervisor/api/auth.py

@@ -11,6 +11,7 @@ from ..addons.addon import Addon
 from ..const import ATTR_PASSWORD, ATTR_USERNAME, REQUEST_FROM
 from ..coresys import CoreSysAttributes
 from ..exceptions import APIForbidden
+from ..utils.json import json_loads
 from .const import CONTENT_TYPE_JSON, CONTENT_TYPE_URL
 from .utils import api_process, api_validate
 
@@ -67,7 +68,7 @@ class APIAuth(CoreSysAttributes):
 
         # Json
         if request.headers.get(CONTENT_TYPE) == CONTENT_TYPE_JSON:
-            data = await request.json()
+            data = await request.json(loads=json_loads)
             return await self._process_dict(request, addon, data)
 
         # URL encoded

supervisor/api/utils.py

@@ -22,7 +22,7 @@ from ..const import (
 from ..coresys import CoreSys
 from ..exceptions import APIError, APIForbidden, DockerAPIError, HassioError
 from ..utils import check_exception_chain, get_message_from_exception_chain
-from ..utils.json import JSONEncoder
+from ..utils.json import json_dumps, json_loads as json_loads_util
 from ..utils.log_format import format_message
 from .const import CONTENT_TYPE_BINARY
 
@@ -48,7 +48,7 @@ def json_loads(data: Any) -> dict[str, Any]:
     if not data:
         return {}
     try:
-        return json.loads(data)
+        return json_loads_util(data)
     except json.JSONDecodeError as err:
         raise APIError("Invalid json") from err
 
@@ -130,7 +130,7 @@ def api_return_error(
             JSON_MESSAGE: message or "Unknown error, see supervisor",
         },
         status=400,
-        dumps=lambda x: json.dumps(x, cls=JSONEncoder),
+        dumps=json_dumps,
     )
 
 
@@ -138,7 +138,7 @@ def api_return_ok(data: dict[str, Any] | None = None) -> web.Response:
     """Return an API ok answer."""
     return web.json_response(
         {JSON_RESULT: RESULT_OK, JSON_DATA: data or {}},
-        dumps=lambda x: json.dumps(x, cls=JSONEncoder),
+        dumps=json_dumps,
     )
 
 

supervisor/utils/json.py

@@ -1,40 +1,63 @@
 """Tools file for Supervisor."""
-from datetime import datetime
+from functools import partial
-import json
 import logging
 from pathlib import Path
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from atomicwrites import atomic_write
+import orjson
 
 from ..exceptions import JsonFileError
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
 
-class JSONEncoder(json.JSONEncoder):
+def json_dumps(data: Any) -> str:
-    """JSONEncoder that supports Supervisor objects."""
+    """Dump json string."""
+    return json_bytes(data).decode("utf-8")
 
-    def default(self, o: Any) -> Any:
-        """Convert Supervisor special objects.
 
-        Hand other objects to the original method.
+def json_encoder_default(obj: Any) -> Any:
-        """
+    """Convert Supervisor special objects."""
-        if isinstance(o, datetime):
+    if isinstance(obj, (set, tuple)):
-            return o.isoformat()
+        return list(obj)
-        if isinstance(o, set):
+    if isinstance(obj, float):
-            return list(o)
+        return float(obj)
-        if isinstance(o, Path):
+    if isinstance(obj, Path):
-            return o.as_posix()
+        return obj.as_posix()
+    raise TypeError
 
-        return super().default(o)
+
+if TYPE_CHECKING:
+
+    def json_bytes(obj: Any) -> bytes:
+        """Dump json bytes."""
+
+else:
+    json_bytes = partial(
+        orjson.dumps,  # pylint: disable=no-member
+        option=orjson.OPT_NON_STR_KEYS,  # pylint: disable=no-member
+        default=json_encoder_default,
+    )
+    """Dump json bytes."""
+
+
+# pylint - https://github.com/ijl/orjson/issues/248
+json_loads = orjson.loads  # pylint: disable=no-member
 
 
 def write_json_file(jsonfile: Path, data: Any) -> None:
     """Write a JSON file."""
     try:
         with atomic_write(jsonfile, overwrite=True) as fp:
-            fp.write(json.dumps(data, indent=2, cls=JSONEncoder))
+            fp.write(
+                orjson.dumps(  # pylint: disable=no-member
+                    data,
+                    option=orjson.OPT_INDENT_2  # pylint: disable=no-member
+                    | orjson.OPT_NON_STR_KEYS,  # pylint: disable=no-member
+                    default=json_encoder_default,
+                ).decode("utf-8")
+            )
         jsonfile.chmod(0o600)
     except (OSError, ValueError, TypeError) as err:
         raise JsonFileError(
@@ -45,7 +68,7 @@ def write_json_file(jsonfile: Path, data: Any) -> None:
 def read_json_file(jsonfile: Path) -> Any:
     """Read a JSON file and return a dict."""
     try:
-        return json.loads(jsonfile.read_text())
+        return json_loads(jsonfile.read_bytes())
     except (OSError, ValueError, TypeError, UnicodeDecodeError) as err:
         raise JsonFileError(
             f"Can't read json from {jsonfile!s}: {err!s}", _LOGGER.error

tests/utils/test_json.py

@@ -1,5 +1,8 @@
 """test json."""
-from supervisor.utils.json import write_json_file
+import time
+from typing import NamedTuple
+
+from supervisor.utils.json import json_dumps, read_json_file, write_json_file
 
 
 def test_file_permissions(tmp_path):
@@ -18,3 +21,42 @@ def test_new_file_permissions(tmp_path):
 
     write_json_file(tempfile, {"test": "data"})
     assert oct(tempfile.stat().st_mode)[-3:] == "600"
+
+
+def test_file_round_trip(tmp_path):
+    """Test file permissions."""
+    tempfile = tmp_path / "test.json"
+    write_json_file(tempfile, {"test": "data"})
+    assert tempfile.is_file()
+    assert oct(tempfile.stat().st_mode)[-3:] == "600"
+    assert read_json_file(tempfile) == {"test": "data"}
+
+
+def test_json_dumps_float_subclass() -> None:
+    """Test the json dumps a float subclass."""
+
+    class FloatSubclass(float):
+        """A float subclass."""
+
+    assert json_dumps({"c": FloatSubclass(1.2)}) == '{"c":1.2}'
+
+
+def test_json_dumps_tuple_subclass() -> None:
+    """Test the json dumps a tuple subclass."""
+
+    tt = time.struct_time((1999, 3, 17, 32, 44, 55, 2, 76, 0))
+
+    assert json_dumps(tt) == "[1999,3,17,32,44,55,2,76,0]"
+
+
+def test_json_dumps_named_tuple_subclass() -> None:
+    """Test the json dumps a tuple subclass."""
+
+    class NamedTupleSubclass(NamedTuple):
+        """A NamedTuple subclass."""
+
+        name: str
+
+    nts = NamedTupleSubclass("a")
+
+    assert json_dumps(nts) == '["a"]'