fix: write json outputs atomically

Author: LuisFigueroaG

arnio/io.py

@@ -2353,6 +2353,43 @@ def write_parquet(
     df.to_parquet(path, **kwargs)
 
 
+@contextmanager
+def _atomic_text_writer(
+    path: str,
+    *,
+    encoding: str,
+    errors: str | None = None,
+    newline: str | None = None,
+) -> Iterator[io.TextIOBase]:
+    directory = os.path.dirname(os.path.abspath(path)) or "."
+    basename = os.path.basename(path)
+    fd, tmp_path_str = tempfile.mkstemp(
+        dir=directory,
+        prefix=f".{basename}.",
+        suffix=".tmp",
+        text=True,
+    )
+    tmp_path: str | None = tmp_path_str
+    try:
+        with os.fdopen(
+            fd,
+            "w",
+            encoding=encoding,
+            errors=errors,
+            newline=newline,
+        ) as dst:
+            yield dst
+        assert tmp_path is not None
+        os.replace(tmp_path, path)
+        tmp_path = None
+    finally:
+        if tmp_path is not None:
+            try:
+                os.unlink(tmp_path)
+            except OSError:
+                pass
+
+
 def write_json(
     frame: ArFrame,
     path: str | os.PathLike[str],
@@ -2421,7 +2458,7 @@ def write_json(
 
     data = frame.to_dict(orient=orient)
 
-    with open(path, "w", encoding="utf-8") as f:
+    with _atomic_text_writer(path, encoding="utf-8") as f:
         json.dump(data, f, indent=indent)
 
 
@@ -2450,9 +2487,8 @@ def write_jsonl(
         ) from exc
 
     try:
-        with open(
+        with _atomic_text_writer(
             path,
-            "w",
             encoding=encoding,
             errors=encoding_errors,
             newline="",

tests/test_write_json.py

@@ -4,6 +4,7 @@
 import pytest
 
 import arnio as ar
+import arnio.io as ar_io
 from arnio.frame import ArFrame
 
 
@@ -64,6 +65,26 @@ def test_write_json_indent(sample_frame: ArFrame, tmp_path: pathlib.Path) -> Non
     assert '    "id":' in content
 
 
+def test_write_json_failure_preserves_existing_file(
+    sample_frame: ArFrame, tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    output_file = tmp_path / "output.json"
+    original = '{"status":"complete"}'
+    output_file.write_text(original, encoding="utf-8")
+
+    def fail_after_partial_write(_data, dst, **_kwargs):
+        dst.write('{"partial":')
+        raise RuntimeError("simulated write failure")
+
+    monkeypatch.setattr(ar_io.json, "dump", fail_after_partial_write)
+
+    with pytest.raises(RuntimeError, match="simulated write failure"):
+        ar.write_json(sample_frame, output_file)
+
+    assert output_file.read_text(encoding="utf-8") == original
+    assert not list(tmp_path.glob(f".{output_file.name}.*.tmp"))
+
+
 def test_write_json_invalid_frame(tmp_path: pathlib.Path) -> None:
     with pytest.raises(TypeError, match="frame must be an ArFrame"):
         ar.write_json({"a": [1]}, tmp_path / "out.json")  # type: ignore

tests/test_write_jsonl.py

@@ -5,6 +5,7 @@
 import pytest
 
 import arnio as ar
+import arnio.io as ar_io
 
 
 def test_write_jsonl_normal_frame(tmp_path):
@@ -23,6 +24,30 @@ def test_write_jsonl_normal_frame(tmp_path):
     ]
 
 
+def test_write_jsonl_failure_preserves_existing_file(tmp_path, monkeypatch):
+    frame = ar.from_pandas(pd.DataFrame({"id": [1, 2]}))
+    path = tmp_path / "existing.jsonl"
+    original = '{"id":0}\n'
+    path.write_text(original, encoding="utf-8")
+    calls = 0
+    original_dumps = ar_io.json.dumps
+
+    def wrapped_dumps(row, **kwargs):
+        nonlocal calls
+        calls += 1
+        if calls == 1:
+            return original_dumps(row, **kwargs)
+        raise ValueError("simulated serialization failure")
+
+    monkeypatch.setattr(ar_io.json, "dumps", wrapped_dumps)
+
+    with pytest.raises(ValueError, match="cannot be serialized"):
+        ar.write_jsonl(frame, path)
+
+    assert path.read_text(encoding="utf-8") == original
+    assert not list(tmp_path.glob(f".{path.name}.*.tmp"))
+
+
 def test_write_jsonl_empty_frame(tmp_path):
     frame = ar.from_pandas(pd.DataFrame({"id": [], "name": []}))