Add format_stem, unique_path, and save_json provenance writers

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01PEAxNzhDCaTrntX3o1JqMV

Author: wmayner

pyphi/provenance.py

@@ -11,8 +11,12 @@
 
 import functools
 import importlib.metadata
+import json
 import platform as _platform
+import re
 import subprocess
+from collections.abc import Mapping
+from dataclasses import asdict
 from dataclasses import dataclass
 from dataclasses import replace
 from datetime import UTC
@@ -158,6 +162,113 @@ def display_rows(self) -> list[tuple[str, str]]:
         return rows
 
 
+def _stem_value(value: Any) -> str:
+    """Format a parameter value for use in a filename stem."""
+    return re.sub(r"[^A-Za-z0-9_+-]", "-", str(value).replace(".", "p"))
+
+
+def format_stem(
+    name: str,
+    params: Mapping[str, Any] | None = None,
+    run_label: str | None = None,
+) -> str:
+    """Build a filename stem encoding a script's parameters.
+
+    Joins ``name``, one ``{key}{value}`` segment per ``params`` entry (in
+    insertion order), and ``run_label`` when given, with underscores.
+    Values and the run label are formatted with ``str()``; ``.`` becomes
+    ``p`` (so ``0.7`` → ``0p7`` and the filename keeps a single suffix)
+    and any character outside ``[A-Za-z0-9_+-]`` becomes ``-``. ``name``
+    is used verbatim.
+
+    Examples
+    --------
+    >>> format_stem("study", {"seed": 42, "noise": 0.7}, "pilot")
+    'study_seed42_noise0p7_pilot'
+    """
+    parts = [name]
+    for key, value in (params or {}).items():
+        parts.append(f"{key}{_stem_value(value)}")
+    if run_label:
+        parts.append(_stem_value(run_label))
+    return "_".join(parts)
+
+
+def unique_path(directory: Path | str, stem: str, suffix: str) -> Path:
+    """Return a non-clobbering path: ``stem+suffix``, else ``stem_v2+suffix``, ...
+
+    Creates ``directory`` (with parents) if it does not exist. Never
+    returns a path that already exists, so earlier outputs are never
+    overwritten.
+    """
+    directory = Path(directory)
+    directory.mkdir(parents=True, exist_ok=True)
+    path = directory / f"{stem}{suffix}"
+    version = 2
+    while path.exists():
+        path = directory / f"{stem}_v{version}{suffix}"
+        version += 1
+    return path
+
+
+def _json_default(obj: Any) -> Any:
+    """``json.dumps`` fallback for numpy values and paths."""
+    if isinstance(obj, np.generic):
+        return obj.item()
+    if isinstance(obj, np.ndarray):
+        return obj.tolist()
+    if isinstance(obj, Path):
+        return str(obj)
+    raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable")
+
+
+def _capture_metadata(
+    params: Mapping[str, Any] | None,
+    seed: int | None,
+    note: str | None,
+) -> Provenance:
+    """Capture a :class:`Provenance`, resolving the seed from ``params``."""
+    if seed is None and params is not None and "seed" in params:
+        seed = int(params["seed"])
+    prov = Provenance.capture(seed=seed)
+    if note is not None:
+        prov = replace(prov, note=note)
+    return prov
+
+
+def save_json(
+    data: Any,
+    directory: Path | str,
+    name: str,
+    *,
+    params: Mapping[str, Any] | None = None,
+    run_label: str | None = None,
+    seed: int | None = None,
+    note: str | None = None,
+) -> Path:
+    """Write ``data`` to a self-describing, non-clobbering JSON file.
+
+    The file holds the envelope ``{"provenance": ..., "params": ...,
+    "data": ...}``. The filename encodes ``params`` and ``run_label``
+    (see :func:`format_stem`); an existing file is never overwritten (a
+    ``_v2``/``_v3`` suffix is added instead). The provenance record
+    stores the seed from ``seed`` or, when omitted, from
+    ``params["seed"]``. numpy scalars and arrays in ``data`` are
+    converted to JSON-native values.
+
+    Returns the written path.
+    """
+    prov = _capture_metadata(params, seed, note)
+    path = unique_path(directory, format_stem(name, params, run_label), ".json")
+    envelope = {
+        "provenance": asdict(prov),
+        "params": dict(params or {}),
+        "data": data,
+    }
+    path.write_text(json.dumps(envelope, indent=2, default=_json_default))
+    return path
+
+
 def _set_provenance(result: Any, prov: Provenance) -> None:
     """Assign ``prov`` to ``result.provenance``, working around frozen results."""
     try:
@@ -204,4 +315,11 @@ def with_provenance(self, **fields: Any) -> HasProvenance:
         return self
 
 
-__all__ = ["HasProvenance", "Provenance", "stamp_wall_time"]
+__all__ = [
+    "HasProvenance",
+    "Provenance",
+    "format_stem",
+    "save_json",
+    "stamp_wall_time",
+    "unique_path",
+]

test/test_provenance_writers.py

@@ -0,0 +1,86 @@
+"""Tests for the script-facing provenance writers."""
+
+import json
+
+import numpy as np
+import pytest
+
+from pyphi import provenance
+
+
+class TestFormatStem:
+    def test_params_in_insertion_order(self):
+        assert (
+            provenance.format_stem("study", {"seed": 42, "trials": 60})
+            == "study_seed42_trials60"
+        )
+
+    def test_float_dot_becomes_p(self):
+        assert provenance.format_stem("study", {"noise": 0.7}) == "study_noise0p7"
+
+    def test_unsafe_characters_sanitized(self):
+        assert provenance.format_stem("study", {"tag": "a/b c"}) == "study_taga-b-c"
+
+    def test_run_label_appended(self):
+        assert (
+            provenance.format_stem("study", {"seed": 1}, "post_reduction")
+            == "study_seed1_post_reduction"
+        )
+
+    def test_no_params(self):
+        assert provenance.format_stem("study") == "study"
+
+
+class TestUniquePath:
+    def test_fresh_path(self, tmp_path):
+        assert provenance.unique_path(tmp_path, "run", ".json") == tmp_path / "run.json"
+
+    def test_versions_on_collision(self, tmp_path):
+        (tmp_path / "run.json").touch()
+        assert (
+            provenance.unique_path(tmp_path, "run", ".json") == tmp_path / "run_v2.json"
+        )
+        (tmp_path / "run_v2.json").touch()
+        assert (
+            provenance.unique_path(tmp_path, "run", ".json") == tmp_path / "run_v3.json"
+        )
+
+    def test_creates_directory(self, tmp_path):
+        target = tmp_path / "a" / "b"
+        provenance.unique_path(target, "run", ".json")
+        assert target.is_dir()
+
+
+class TestSaveJson:
+    def test_envelope_and_numpy_values(self, tmp_path):
+        path = provenance.save_json(
+            {"phi": np.float64(0.5), "counts": np.arange(3)},
+            tmp_path,
+            "study",
+            params={"seed": 42},
+        )
+        assert path == tmp_path / "study_seed42.json"
+        document = json.loads(path.read_text())
+        assert document["data"] == {"phi": 0.5, "counts": [0, 1, 2]}
+        assert document["params"] == {"seed": 42}
+        assert document["provenance"]["seed"] == 42
+        assert document["provenance"]["pyphi_version"]
+        assert document["provenance"]["timestamp"]
+
+    def test_explicit_seed_overrides_params(self, tmp_path):
+        path = provenance.save_json({}, tmp_path, "study", params={"seed": 42}, seed=7)
+        assert json.loads(path.read_text())["provenance"]["seed"] == 7
+
+    def test_note_stored(self, tmp_path):
+        path = provenance.save_json({}, tmp_path, "study", note="pilot run")
+        assert json.loads(path.read_text())["provenance"]["note"] == "pilot run"
+
+    def test_unserializable_payload_raises(self, tmp_path):
+        with pytest.raises(TypeError, match="not JSON serializable"):
+            provenance.save_json({"bad": object()}, tmp_path, "study")
+
+    def test_no_clobber(self, tmp_path):
+        first = provenance.save_json({"x": 1}, tmp_path, "study", params={"seed": 1})
+        second = provenance.save_json({"x": 2}, tmp_path, "study", params={"seed": 1})
+        assert second == tmp_path / "study_seed1_v2.json"
+        assert json.loads(first.read_text())["data"] == {"x": 1}