perf: optionally use 'orjson' if available

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"

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 = []

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("")

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
 

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,

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:

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

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
 

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

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
 

src/taskgraph/util/json.py

@@ -0,0 +1,73 @@
+# 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")
+
+    separators = (",", ": ") if indent else (",", ":")
+    return json.dumps(
+        obj,
+        default=default,
+        indent=indent,
+        sort_keys=sort_keys,
+        separators=separators,
+    )
+
+
+def dump(obj: Any, fh: SupportsWrite[str], **kwargs) -> None:
+    fh.write(dumps(obj, **kwargs))

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

taskcluster/kinds/check/kind.yml

@@ -38,4 +38,4 @@ tasks:
             symbol: check(type)
         run:
             command: >-
-                uv run pyright
+                uv run --all-extras pyright

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",

test/test_util_json.py

@@ -0,0 +1,76 @@
+from itertools import count
+from textwrap import dedent
+from unittest.mock import Mock
+
+import pytest
+
+from taskgraph.util.json import dump, dumps, load, loads
+
+
+@pytest.fixture(autouse=True, params=["json", "orjson"])
+def with_module(request, mocker):
+    module = request.param
+    if module == "orjson":
+        pytest.importorskip("orjson")
+
+    if module == "json":
+        mocker.patch("taskgraph.util.json.orjson", None)
+
+
+@pytest.mark.parametrize(
+    "input_data, expected",
+    [
+        ('{"key": "value"}', {"key": "value"}),
+    ],
+    ids=count(),
+)
+def test_loads(input_data, expected):
+    assert loads(input_data) == expected
+
+
+@pytest.mark.parametrize(
+    "input_data, expected",
+    [
+        ('{"key": "value"}', {"key": "value"}),
+    ],
+    ids=count(),
+)
+def test_load(input_data, expected):
+    mock_file = Mock(read=Mock(return_value=input_data))
+    assert load(mock_file) == expected
+
+
+@pytest.mark.parametrize(
+    "input_data, json_args, expected",
+    [
+        ({"b": 1, "a": 2}, {}, '{"b":1,"a":2}'),
+        ({"b": 1, "a": 2}, {"sort_keys": True}, '{"a":2,"b":1}'),
+        (
+            {"b": 1, "a": 2},
+            {"indent": 2},
+            dedent("""
+            {
+              "b": 1,
+              "a": 2
+            }""").lstrip(),
+        ),
+    ],
+    ids=count(),
+)
+def test_dumps(input_data, json_args, expected):
+    assert dumps(input_data, **json_args) == expected
+
+
+@pytest.mark.parametrize(
+    "input_data, json_args, expected",
+    [
+        ({"key": "value"}, {}, '{"key":"value"}'),
+    ],
+    ids=count(),
+)
+def test_dump(input_data, json_args, expected):
+    mock_file = Mock()
+
+    dump(input_data, mock_file, **json_args)
+
+    mock_file.write.assert_called_once_with(expected)