feat: preserve tuples in JSON

tests/unit/cli/test_cli.py

@@ -1,10 +1,9 @@
 # pylint: disable=unused-argument,redefined-outer-name
-import json
-
 import pytest
 from click.testing import CliRunner
 
 from zetta_utils import builder, cli
+from zetta_utils.parsing import json
 
 
 @pytest.fixture

tests/unit/parsing/test_json.py

@@ -0,0 +1,36 @@
+from io import StringIO
+
+from zetta_utils.parsing import json
+
+
+def test_str_roundtrip():
+    assert json.loads(json.dumps([1, "2"])) == [1, "2"]
+
+
+def test_str_tuple_roundtrip():
+    assert json.loads(json.dumps((1, "2"))) == (1, "2")
+
+
+def test_str_nested_tuple_roundtrip():
+    assert json.loads(json.dumps((1, "2", (3, "4")))) == (1, "2", (3, "4"))
+
+
+def test_fp_roundtrip():
+    fp = StringIO()
+    json.dump([1, "2"], fp)
+    fp.seek(0)
+    assert json.load(fp) == [1, "2"]
+
+
+def test_fp_tuple_roundtrip():
+    fp = StringIO()
+    json.dump((1, "2"), fp)
+    fp.seek(0)
+    assert json.load(fp) == (1, "2")
+
+
+def test_fp_nested_tuple_roundtrip():
+    fp = StringIO()
+    json.dump((1, "2", (3, "4")), fp)
+    fp.seek(0)
+    assert json.load(fp) == (1, "2", (3, "4"))

zetta_utils/builder/build.py

@@ -1,14 +1,14 @@
 """Bulding objects from nested specs."""
 from __future__ import annotations
 
-import json
 from typing import Any, Callable, Final, Optional
 
 import attrs
 from typeguard import typechecked
 
 from zetta_utils import parsing
 from zetta_utils.common import ctx_managers
+from zetta_utils.parsing import json
 
 from . import constants
 from .registry import get_matching_entry
@@ -183,7 +183,6 @@ def get_display_name(self):  # pragma: no cover # pretty print
 
     def _get_built_spec_kwargs(self, version: str) -> dict[str, Any]:
         if self._built_spec_kwargs is None:
-
             self._built_spec_kwargs = {
                 k: _traverse_spec(
                     v,

zetta_utils/cli/main.py

@@ -1,4 +1,3 @@
-import json
 import os
 import pprint
 from typing import Optional
@@ -7,6 +6,7 @@
 
 import zetta_utils
 from zetta_utils import log
+from zetta_utils.parsing import json
 
 logger = log.get_logger("zetta_utils")
 

zetta_utils/cloud_management/execution_tracker.py

@@ -1,4 +1,3 @@
-import json
 import os
 import time
 from contextlib import contextmanager
@@ -14,6 +13,7 @@
 from zetta_utils.layer.db_layer import DBRowDataT, build_db_layer
 from zetta_utils.layer.db_layer.datastore import DatastoreBackend
 from zetta_utils.log import get_logger
+from zetta_utils.parsing import json
 
 from .resource_allocation.k8s import ClusterInfo
 

zetta_utils/parsing/__init__.py

@@ -1,2 +1,3 @@
 from . import cue
 from . import ngl_state
+from . import json

zetta_utils/parsing/cue.py

@@ -1,5 +1,4 @@
 """cuelang parsing."""
-import json
 import os
 import pathlib
 import subprocess
@@ -8,6 +7,7 @@
 import fsspec
 
 from zetta_utils import log
+from zetta_utils.parsing import json
 
 logger = log.get_logger("zetta_utils")
 

zetta_utils/parsing/json.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+import json
+from collections.abc import Iterator
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from _typeshed import SupportsRead, SupportsWrite
+
+
+def _mark_python_types(obj: Any) -> Any:
+    if isinstance(obj, tuple):
+        return {"__tuple__": [_mark_python_types(e) for e in obj]}
+    if isinstance(obj, list):
+        return [_mark_python_types(e) for e in obj]
+    if isinstance(obj, dict):
+        return {key: _mark_python_types(value) for key, value in obj.items()}
+    else:
+        return obj
+
+
+class ZettaSpecJSONEncoder(json.JSONEncoder):
+    def encode(self, o: Any) -> str:
+        return super().encode(_mark_python_types(o))
+
+    def iterencode(self, o: Any, _one_shot: bool = False) -> Iterator[str]:
+        return super().iterencode(_mark_python_types(o), _one_shot=_one_shot)
+
+
+def tuple_hook(obj):
+    if "__tuple__" in obj:
+        return tuple(obj["__tuple__"])
+    else:
+        return obj
+
+
+def dumps(obj, **kwargs) -> str:
+    return json.dumps(obj, cls=ZettaSpecJSONEncoder, **kwargs)
+
+
+def dump(obj: Any, fp: SupportsWrite[str], **kwargs) -> None:
+    json.dump(obj, fp, cls=ZettaSpecJSONEncoder, **kwargs)
+
+
+def loads(s: str, **kwargs) -> Any:
+    return json.loads(s, object_hook=tuple_hook, **kwargs)
+
+
+def load(fp: SupportsRead[str | bytes], **kwargs) -> Any:
+    return json.load(fp, object_hook=tuple_hook, **kwargs)

zetta_utils/training/lightning/train.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import json
 import os
 from contextlib import ExitStack
 from typing import Any, Dict, Final, List, Optional
@@ -16,6 +15,7 @@
 from zetta_utils import builder, load_all_modules, log, mazepa, parsing
 from zetta_utils.builder.build import BuilderPartial
 from zetta_utils.cloud_management import execution_tracker, resource_allocation
+from zetta_utils.parsing import json
 
 logger = log.get_logger("zetta_utils")
 

zetta_utils/training/lightning/trainers/default.py

@@ -2,7 +2,6 @@
 
 import datetime
 import importlib.metadata
-import json
 import os
 from typing import Any, Dict, List, Optional
 
@@ -15,6 +14,7 @@
 from pytorch_lightning.strategies import ddp
 
 from zetta_utils import builder, log
+from zetta_utils.parsing import json
 
 logger = log.get_logger("zetta_utils")
 ONNX_OPSET_VERSION = 17
@@ -126,7 +126,6 @@ def log_config(config):
     def save_checkpoint(
         self, filepath, weights_only: bool = False, storage_options: Optional[Any] = None
     ):  # pylint: disable=too-many-locals
-
         if filepath.startswith("./"):
             filepath = f"{self.default_root_dir}/{filepath[2:]}"
         super().save_checkpoint(filepath, weights_only, storage_options)