Write TSV, JSON and text files atomically

save_json(), write_python(), write_text(), write_tsv() and _write_tsv_simple()
all opened the destination in 'w' mode, which truncates it before a single byte
of the new content is written. A crash, a full disk or a kill in that window
leaves the file empty or half written and the previous contents are gone.

cluster_group.tsv is the one that hurts, since it holds manual curation labels
that cannot be recomputed.

Add _atomic_open(), a context manager that writes to a temporary file in the
destination directory and moves it into place with os.replace() on a clean
exit, and route all five writers through it. The temporary file is given the
permissions of the file it replaces, or the permissions open() would have used
for a new file, since tempfile creates its files with 0600.

Author: adityasingh2400

CHANGELOG.md

@@ -14,6 +14,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Development checkouts now report a development version with a Git commit suffix.
 - Template datasets without waveform templates can be reloaded after cluster assignments change,
   and stored spike-waveform subsets provide their waveform sample count.
+- TSV, JSON, text and `params.py` files are written atomically, so a crash during a save no longer
+  truncates the file it was replacing, for instance `cluster_group.tsv`.
 
 ## [2.7.0] 2025-12-10
 

phylib/utils/_misc.py

@@ -8,13 +8,15 @@
 #------------------------------------------------------------------------------
 
 import base64
+from contextlib import contextmanager
 import csv
 from importlib import import_module
 import json
 import logging
 import os
 from pathlib import Path
 import subprocess
+import tempfile
 from textwrap import dedent
 
 import numpy as np
@@ -24,6 +26,54 @@
 logger = logging.getLogger(__name__)
 
 
+#------------------------------------------------------------------------------
+# Atomic file writing
+#------------------------------------------------------------------------------
+
+def _copy_file_mode(src, dst):
+    """Give `src` the permissions `dst` has, or the permissions a new file would get."""
+    try:
+        mode = os.stat(dst).st_mode & 0o777
+    except OSError:
+        # The destination does not exist yet: reproduce what open() would have done, since
+        # tempfile creates its files with 0600 and saving must not make the file private.
+        umask = os.umask(0)
+        os.umask(umask)
+        mode = 0o666 & ~umask
+    os.chmod(src, mode)
+
+
+@contextmanager
+def _atomic_open(path, mode='w', **kwargs):
+    """Open a file for writing without ever truncating the file it replaces.
+
+    The contents are written to a temporary file in the same directory as the destination, and
+    that file is moved into place with `os.replace()` once the `with` block completes, which is
+    atomic as long as both paths are on the same filesystem. A reader therefore sees either the
+    old file or the new one, never a half-written one. If the block raises, the temporary file
+    is removed and any existing file is left untouched.
+
+    """
+    path = Path(path)
+    ensure_dir_exists(path.parent)
+    # NOTE: the temporary file is hidden and does not end in a data file extension, so that a
+    # file left behind by a hard kill is not picked up by the globs used to discover datasets.
+    fd, tmp = tempfile.mkstemp(dir=str(path.parent), prefix='.' + path.name + '.', suffix='.tmp')
+    try:
+        with os.fdopen(fd, mode, **kwargs) as f:
+            yield f
+            f.flush()
+            os.fsync(f.fileno())
+        _copy_file_mode(tmp, path)
+        os.replace(tmp, path)
+    except BaseException:
+        try:
+            os.remove(tmp)
+        except OSError:  # pragma: no cover
+            pass
+        raise
+
+
 #------------------------------------------------------------------------------
 # JSON utility functions
 #------------------------------------------------------------------------------
@@ -128,9 +178,7 @@ def save_json(path, data):
     """
     assert isinstance(data, dict)
     data = _stringify_keys(data)
-    path = Path(path)
-    ensure_dir_exists(path.parent)
-    with path.open('w') as f:
+    with _atomic_open(path) as f:
         json.dump(data, f, cls=_CustomEncoder, indent=2, sort_keys=True)
 
 
@@ -190,7 +238,7 @@ def write_python(path, data):
     -------
 
     """
-    with open(path, 'w') as f:
+    with _atomic_open(path) as f:
         for k, v in data.items():
             if isinstance(v, str):
                 v = '"%s"' % v
@@ -206,9 +254,8 @@ def read_text(path):
 def write_text(path, contents):
     """Write a text file."""
     contents = dedent(contents)
-    path = Path(path)
-    ensure_dir_exists(path.parent)
-    path.write_text(contents)
+    with _atomic_open(path) as f:
+        f.write(contents)
 
 
 def _try_make_number(value):
@@ -266,9 +313,8 @@ def write_tsv(path, data, first_field=None, exclude_fields=(), n_significant_fig
 
     """
     path = Path(path)
-    ensure_dir_exists(path.parent)
     delimiter = '\t' if path.suffix == '.tsv' else ','
-    with path.open('w', newline='') as f:
+    with _atomic_open(path, newline='') as f:
         if not data:
             logger.info("Data was empty when writing %s.", path)
             return
@@ -327,9 +373,8 @@ def _write_tsv_simple(path, field_name, data):
 
     """
     path = Path(path)
-    ensure_dir_exists(path.parent)
     delimiter = '\t' if path.suffix == '.tsv' else ','
-    with path.open('w', newline='') as f:
+    with _atomic_open(path, newline='') as f:
         writer = csv.writer(f, delimiter=delimiter)
         writer.writerow(['cluster_id', field_name])
         writer.writerows([(cluster_id, data[cluster_id]) for cluster_id in sorted(data)])

phylib/utils/tests/test_misc.py

@@ -6,15 +6,19 @@
 # Imports
 #------------------------------------------------------------------------------
 
+import csv
+import json
+import os
 import re
+import stat
 
 import numpy as np
 from numpy.testing import assert_array_equal as ae
 from pytest import raises, mark
 
 from .._misc import (
-    _git_version, load_json, save_json, load_pickle, save_pickle, read_python, write_python,
-    read_text, write_text, _read_tsv_simple, _write_tsv_simple, read_tsv, write_tsv,
+    _atomic_open, _git_version, load_json, save_json, load_pickle, save_pickle, read_python,
+    write_python, read_text, write_text, _read_tsv_simple, _write_tsv_simple, read_tsv, write_tsv,
     _pretty_floats, _encode_qbytearray, _decode_qbytearray, _fullname, _load_from_fullname)
 
 
@@ -137,6 +141,91 @@ def test_write_tsv(tempdir):
     assert read_tsv(path) == data
 
 
+class _Boom(Exception):
+    pass
+
+
+def test_atomic_open(tempdir):
+    path = tempdir / 'test_dir/test.txt'
+
+    with _atomic_open(path) as f:
+        f.write('hello')
+    assert read_text(path) == 'hello'
+    # The temporary file must not be left behind.
+    assert sorted(p.name for p in path.parent.iterdir()) == ['test.txt']
+
+    # A new file gets the permissions open() would have given it, not the 0600 of a
+    # temporary file.
+    umask = os.umask(0)
+    os.umask(umask)
+    assert stat.S_IMODE(path.stat().st_mode) == 0o666 & ~umask
+
+    # The permissions of an existing file are preserved.
+    path.chmod(0o640)
+    with _atomic_open(path) as f:
+        f.write('world')
+    assert read_text(path) == 'world'
+    assert stat.S_IMODE(path.stat().st_mode) == 0o640
+
+    # A failure half-way through leaves the previous file untouched and no leftovers.
+    with raises(_Boom):
+        with _atomic_open(path) as f:
+            f.write('this should never be visible')
+            raise _Boom()
+    assert read_text(path) == 'world'
+    assert sorted(p.name for p in path.parent.iterdir()) == ['test.txt']
+
+
+@mark.parametrize('writer', ['write_tsv', '_write_tsv_simple'])
+def test_write_tsv_atomic(tempdir, monkeypatch, writer):
+    # Curation data such as cluster_group.tsv must survive a crash during a save.
+    path = tempdir / 'cluster_group.tsv'
+    if writer == 'write_tsv':
+        def _write(value):
+            write_tsv(path, [{'cluster_id': 2, 'group': value}], first_field='cluster_id')
+    else:
+        def _write(value):
+            _write_tsv_simple(path, 'group', {2: value})
+
+    _write('good')
+    before = path.read_bytes()
+    assert before
+    files_before = sorted(p.name for p in tempdir.iterdir())
+
+    real_writer = csv.writer
+
+    def _failing_writer(f, **kwargs):
+        # Mimic a crash after part of the file has been written.
+        real_writer(f, **kwargs).writerow(['cluster_id', 'group'])
+        raise _Boom()
+
+    monkeypatch.setattr(csv, 'writer', _failing_writer)
+    with raises(_Boom):
+        _write('mua')
+
+    assert path.read_bytes() == before
+    assert sorted(p.name for p in tempdir.iterdir()) == files_before
+
+
+def test_save_json_atomic(tempdir, monkeypatch):
+    path = tempdir / 'test.json'
+    save_json(path, {'a': 1})
+    before = path.read_bytes()
+    assert before
+    files_before = sorted(p.name for p in tempdir.iterdir())
+
+    def _failing_dump(data, f, **kwargs):
+        f.write('{"a":')
+        raise _Boom()
+
+    monkeypatch.setattr(json, 'dump', _failing_dump)
+    with raises(_Boom):
+        save_json(path, {'a': 2})
+
+    assert path.read_bytes() == before
+    assert sorted(p.name for p in tempdir.iterdir()) == files_before
+
+
 def test_git_version():
     v = _git_version()
     assert re.match(r'^\+git\.[0-9a-f]{8}(?:\.dirty)?$', v)