agentic_upsampling/io_utils.py

"""Small JSON and file helpers for agentic upsampling runs."""

from __future__ import annotations

import json
import os
import tempfile
from pathlib import Path
from typing import Any


def write_json_atomic(path: Path, data: Any, *, ensure_ascii: bool = True) -> None:
    """Write JSON through a temporary file and atomically replace the destination."""
    path.parent.mkdir(parents=True, exist_ok=True)
    fd, tmp_name = tempfile.mkstemp(prefix=f".{path.name}.", suffix=".tmp", dir=path.parent)
    try:
        with os.fdopen(fd, "w", encoding="utf-8") as f:
            json.dump(data, f, ensure_ascii=ensure_ascii, indent=2)
            f.write("\n")
        Path(tmp_name).replace(path)
    except Exception:
        try:
            Path(tmp_name).unlink(missing_ok=True)
        finally:
            raise


def append_jsonl(path: Path, data: Any, *, ensure_ascii: bool = True) -> None:
    """Append one compact JSON record to a JSONL file."""
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("a", encoding="utf-8") as f:
        f.write(json.dumps(data, ensure_ascii=ensure_ascii, separators=(",", ":")) + "\n")


def read_json(path: Path) -> dict[str, Any]:
    """Read a JSON object from disk."""
    data = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(data, dict):
        raise ValueError(f"{path} must contain a JSON object.")
    return data


def compact_json(data: dict[str, Any], *, ensure_ascii: bool = True) -> str:
    """Serialize JSON using the compact prompt format expected by the generation endpoint."""
    return json.dumps(data, ensure_ascii=ensure_ascii, separators=(",", ":"))