perf: optionally use 'orjson' if available

ahal · ahal · commit d5fd00d76ef1 · 2025-05-15T16:38:06.000-04:00
diff --git a/pyproject.toml b/pyproject.toml
@@ -34,9 +34,8 @@ dependencies = [
 dynamic = ["version"]  # set in src/taskgraph/__init__.py
 
 [project.optional-dependencies]
-load-image = [
-  "zstandard"
-]
+load-image = ["zstandard"]
+orjson = ["orjson"]
 
 [project.scripts]
 taskgraph = "taskgraph.main:main"
diff --git a/src/taskgraph/actions/registry.py b/src/taskgraph/actions/registry.py
@@ -4,7 +4,6 @@
 
 
 import functools
-import json
 from collections import namedtuple
 from types import FunctionType
 
@@ -13,7 +12,7 @@
 from taskgraph import create
 from taskgraph.config import load_graph_config
 from taskgraph.parameters import Parameters
-from taskgraph.util import hash, taskcluster, yaml
+from taskgraph.util import hash, json, taskcluster, yaml
 from taskgraph.util.python_path import import_sibling_modules
 
 actions = []
diff --git a/src/taskgraph/create.py b/src/taskgraph/create.py
@@ -3,14 +3,14 @@
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
 
-import json
 import logging
 import os
 import sys
 from concurrent import futures
 
 from slugid import nice as slugid
 
+from taskgraph.util import json
 from taskgraph.util.parameterization import resolve_timestamps
 from taskgraph.util.taskcluster import CONCURRENCY, get_session
 from taskgraph.util.time import current_json_time
@@ -116,8 +116,7 @@ def create_task(session, task_id, label, task_def):
             [task_id, task_def],
             sys.stdout,
             sort_keys=True,
-            indent=4,
-            separators=(",", ": "),
+            indent=2,
         )
         # add a newline
         print("")
diff --git a/src/taskgraph/decision.py b/src/taskgraph/decision.py
@@ -3,7 +3,6 @@
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
 
-import json
 import logging
 import os
 import pathlib
@@ -19,6 +18,7 @@
 from taskgraph.generator import TaskGraphGenerator
 from taskgraph.parameters import Parameters, get_version
 from taskgraph.taskgraph import TaskGraph
+from taskgraph.util import json
 from taskgraph.util.python_path import find_object
 from taskgraph.util.schema import Schema, validate_schema
 from taskgraph.util.vcs import Repository, get_repository
@@ -374,7 +374,7 @@ def write_artifact(filename, data):
             yaml.safe_dump(data, f, allow_unicode=True, default_flow_style=False)
     elif filename.endswith(".json"):
         with open(path, "w") as f:
-            json.dump(data, f, sort_keys=True, indent=2, separators=(",", ": "))
+            json.dump(data, f, sort_keys=True, indent=2)
     elif filename.endswith(".gz"):
         import gzip
 
diff --git a/src/taskgraph/docker.py b/src/taskgraph/docker.py
@@ -3,7 +3,6 @@
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
 
-import json
 import os
 import shlex
 import subprocess
@@ -18,7 +17,7 @@
 except ImportError as e:
     zstd = e
 
-from taskgraph.util import docker
+from taskgraph.util import docker, json
 from taskgraph.util.taskcluster import (
     get_artifact_url,
     get_root_url,
diff --git a/src/taskgraph/main.py b/src/taskgraph/main.py
@@ -4,7 +4,6 @@
 
 import argparse
 import atexit
-import json
 import logging
 import os
 import re
@@ -56,9 +55,9 @@ def format_taskgraph_labels(taskgraph):
 
 
 def format_taskgraph_json(taskgraph):
-    return json.dumps(
-        taskgraph.to_json(), sort_keys=True, indent=2, separators=(",", ": ")
-    )
+    from taskgraph.util import json
+
+    return json.dumps(taskgraph.to_json(), sort_keys=True, indent=2)
 
 
 def format_taskgraph_yaml(taskgraph):
@@ -757,6 +756,7 @@ def actions(args):
     from taskgraph.actions import render_actions_json
     from taskgraph.generator import TaskGraphGenerator
     from taskgraph.parameters import parameters_loader
+    from taskgraph.util import json
 
     if args.pop("verbose", False):
         logging.root.setLevel(logging.DEBUG)
@@ -766,7 +766,7 @@ def actions(args):
         tgg = TaskGraphGenerator(root_dir=args.get("root"), parameters=parameters)
 
         actions = render_actions_json(tgg.parameters, tgg.graph_config, "DECISION-TASK")
-        print(json.dumps(actions, sort_keys=True, indent=2, separators=(",", ": ")))
+        print(json.dumps(actions, sort_keys=True, indent=2))
     except Exception:
         traceback.print_exc()
         sys.exit(1)
@@ -784,6 +784,7 @@ def actions(args):
 def action_callback(options):
     from taskgraph.actions import trigger_action_callback
     from taskgraph.actions.util import get_parameters
+    from taskgraph.util import json
 
     try:
         # the target task for this action (or null if it's a group action)
@@ -833,7 +834,7 @@ def test_action_callback(options):
     import taskgraph.actions
     import taskgraph.parameters
     from taskgraph.config import load_graph_config
-    from taskgraph.util import yaml
+    from taskgraph.util import json, yaml
 
     def load_data(filename):
         with open(filename) as f:
diff --git a/src/taskgraph/parameters.py b/src/taskgraph/parameters.py
@@ -4,7 +4,6 @@
 
 import gzip
 import hashlib
-import json
 import os
 import time
 from datetime import datetime
@@ -18,7 +17,7 @@
 import mozilla_repo_urls
 from voluptuous import ALLOW_EXTRA, Any, Optional, Required, Schema
 
-from taskgraph.util import yaml
+from taskgraph.util import json, yaml
 from taskgraph.util.readonlydict import ReadOnlyDict
 from taskgraph.util.schema import validate_schema
 from taskgraph.util.taskcluster import find_task_id, get_artifact_url
diff --git a/src/taskgraph/transforms/docker_image.py b/src/taskgraph/transforms/docker_image.py
@@ -2,8 +2,6 @@
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
-
-import json
 import logging
 import os
 import re
@@ -12,6 +10,7 @@
 
 import taskgraph
 from taskgraph.transforms.base import TransformSequence
+from taskgraph.util import json
 from taskgraph.util.docker import create_context_tar, generate_context_hash
 from taskgraph.util.schema import Schema
 
diff --git a/src/taskgraph/transforms/run/__init__.py b/src/taskgraph/transforms/run/__init__.py
@@ -10,14 +10,14 @@
 """
 
 import copy
-import json
 import logging
 
 from voluptuous import Exclusive, Extra, Optional, Required
 
 from taskgraph.transforms.base import TransformSequence
 from taskgraph.transforms.cached_tasks import order_tasks
 from taskgraph.transforms.task import task_description_schema
+from taskgraph.util import json
 from taskgraph.util import path as mozpath
 from taskgraph.util.python_path import import_sibling_modules
 from taskgraph.util.schema import Schema, validate_schema
diff --git a/src/taskgraph/transforms/run/common.py b/src/taskgraph/transforms/run/common.py
@@ -7,11 +7,10 @@
 consistency.
 """
 
-import json
 from typing import Any, Dict, List, Union
 
 from taskgraph.transforms.base import TransformConfig
-from taskgraph.util import path
+from taskgraph.util import json, path
 from taskgraph.util.caches import CACHES, get_checkout_dir
 from taskgraph.util.taskcluster import get_artifact_prefix
 
diff --git a/src/taskgraph/util/json.py b/src/taskgraph/util/json.py
@@ -0,0 +1,72 @@
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+"""
+This module is a compatibility shim between stdlib json and orjson.
+It uses orjson if available, and falls back to stdlib json.
+
+It only supports features and parameters that both modules support,
+and attempts to produce identical output regardless of which module
+is used.
+
+The interface is a subset of what's supported by stdlib json.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import TYPE_CHECKING, Any, Callable
+
+if TYPE_CHECKING:
+    from _typeshed import SupportsRead, SupportsWrite
+
+try:
+    import orjson
+except ImportError:
+    orjson = None
+
+
+def loads(s: str | bytes | bytearray) -> Any:
+    if orjson:
+        return orjson.loads(s)
+
+    return json.loads(s)
+
+
+def load(fh: SupportsRead[str | bytes]) -> Any:
+    if orjson:
+        return loads(fh.read())
+    return json.load(fh)
+
+
+def dumps(
+    obj: Any,
+    default: Callable[[Any], Any] | None = None,
+    indent: int | None = None,
+    sort_keys: bool = False,
+) -> str:
+    if indent and indent != 2:
+        raise ValueError("An indent other than 2 is not supported!")
+
+    if orjson:
+        option = 0
+        if indent:
+            option |= orjson.OPT_INDENT_2
+
+        if sort_keys:
+            option |= orjson.OPT_SORT_KEYS
+
+        return orjson.dumps(obj, default=default, option=option).decode("utf-8")
+
+    return json.dumps(
+        obj,
+        default=default,
+        indent=indent,
+        sort_keys=sort_keys,
+        separators=(",", ":"),
+    )
+
+
+def dump(obj: Any, fh: SupportsWrite[str], **kwargs) -> None:
+    fh.write(dumps(obj, **kwargs))
diff --git a/taskcluster/docker/decision/system-setup.sh b/taskcluster/docker/decision/system-setup.sh
@@ -9,7 +9,7 @@ apt-get install -y --force-yes --no-install-recommends \
     python3-pip
 
 pushd /setup/taskgraph
-uv export --no-cache --no-dev > /setup/requirements.txt
+uv export --no-cache --extra orjson --no-dev > /setup/requirements.txt
 uv pip install --no-cache --system --break-system-packages -r /setup/requirements.txt
 uv pip install --no-cache --system --break-system-packages --no-deps .
 popd
diff --git a/taskcluster/kinds/check/kind.yml b/taskcluster/kinds/check/kind.yml
@@ -38,4 +38,4 @@ tasks:
             symbol: check(type)
         run:
             command: >-
-                uv run pyright
+                uv run --all-extras pyright
diff --git a/test/test_actions_rebuild_cached_tasks.py b/test/test_actions_rebuild_cached_tasks.py
@@ -51,5 +51,5 @@ def test_rebuild_cached_tasks(run_action):
         make_task(label="bar", task_def={"name": "bar"}),
     )
     out, _ = run_action("rebuild-cached-tasks", graph)
-    assert '"name": "foo"' in out
-    assert '"name": "bar"' not in out
+    assert '"name":"foo"' in out
+    assert '"name":"bar"' not in out
diff --git a/test/test_transforms_run_run_task.py b/test/test_transforms_run_run_task.py
@@ -80,7 +80,7 @@ def assert_docker_worker(task):
                 "CI_REPOSITORY_TYPE": "hg",
                 "HG_STORE_PATH": "/builds/worker/checkouts/hg-store",
                 "MOZ_SCM_LEVEL": "1",
-                "REPOSITORIES": '{"ci": "Taskgraph"}',
+                "REPOSITORIES": '{"ci":"Taskgraph"}',
                 "VCS_PATH": "/builds/worker/checkouts/vcs",
             },
             "implementation": "docker-worker",
@@ -116,7 +116,7 @@ def assert_generic_worker(task):
                 "CI_REPOSITORY_TYPE": "hg",
                 "HG_STORE_PATH": "y:/hg-shared",
                 "MOZ_SCM_LEVEL": "1",
-                "REPOSITORIES": '{"ci": "Taskgraph"}',
+                "REPOSITORIES": '{"ci":"Taskgraph"}',
                 "VCS_PATH": "build/src",
             },
             "implementation": "generic-worker",
diff --git a/uv.lock b/uv.lock
