feat: atomic_write helper (#157)

dol/__init__.py

@@ -114,6 +114,8 @@ def ihead(store, n=1):
     subfolder_stores,  # a store of stores, each store corresponding to a subfolder
 )
 
+from dol.atomic import atomic_write  # crash-safe file write (temp file + os.replace)
+
 from dol.util import (
     non_colliding_key,  # make a key that does not collide with existing keys
     get_app_folder,  # get the path to a folder for application data

dol/atomic.py

@@ -0,0 +1,118 @@
+"""Atomic file writes: write to a sibling temp file, then ``os.replace`` it into place.
+
+A crash (or an exception while producing the data) leaves the old content or the new,
+never a torn or empty file. Dependency-free; ``atomic_write`` is the one entry point.
+
+    >>> import os, tempfile
+    >>> path = os.path.join(tempfile.mkdtemp(), "f.txt")
+    >>> atomic_write(path, b"hello")
+    >>> open(path).read(), os.listdir(os.path.dirname(path))
+    ('hello', ['f.txt'])
+    >>> atomic_write(path, (c.encode() for c in "ab"))  # chunks are streamed
+    >>> open(path).read()
+    'ab'
+"""
+
+import os
+import stat
+import uuid
+from collections.abc import Iterable
+
+_BYTES_LIKE = (bytes, bytearray, memoryview)
+
+
+def _fsync_dir(dirpath: str) -> None:
+    """Best-effort fsync of a directory, so a rename survives power loss (POSIX only)."""
+    try:
+        fd = os.open(dirpath or ".", os.O_RDONLY)
+    except OSError:  # e.g. Windows cannot open directories
+        return
+    try:
+        os.fsync(fd)
+    except OSError:
+        pass
+    finally:
+        os.close(fd)
+
+
+def _mode_to_apply(path: str, perms: "int | None") -> "int | None":
+    """``perms`` if given, else the mode of the file being replaced (``None`` if new)."""
+    if perms is not None:
+        return perms
+    try:
+        return stat.S_IMODE(os.stat(path).st_mode)
+    except OSError:
+        return None
+
+
+def atomic_write(
+    path,
+    data: "bytes | Iterable[bytes]",
+    *,
+    fsync: bool = False,
+    perms: "int | None" = None,
+) -> None:
+    """Write ``data`` to ``path`` atomically (temp file in the same folder, then ``os.replace``).
+
+    :param path: Target file; its parent folder must exist. A symlink is followed:
+        its target is what gets replaced.
+    :param data: Bytes-like, or an iterable of bytes-like chunks (written one at a
+        time, so a large payload need not be held in memory). If the iterable raises,
+        ``path`` is untouched. Text must be encoded first.
+    :param fsync: Flush the file (and, on POSIX, the folder) to disk before returning:
+        durable across power loss, not just across a process crash. Slower. (On macOS
+        ``os.fsync`` is not ``F_FULLFSYNC``: best effort.)
+    :param perms: Mode (e.g. ``0o600``) for the file, set at creation so the data is
+        never readable more widely, even briefly, and not narrowed by the umask.
+        Default: the mode of the file being replaced, else the usual umask mode.
+
+    The temp file is a hidden, uniquely named sibling (``.<name>.<random>.tmp``), created
+    with ``O_EXCL``, so concurrent writers never share one and key listings skip it.
+    On any failure it is removed and the old content is untouched. Atomic on POSIX,
+    and on Windows within one volume. Hard links to the old file are detached, and its owner is not preserved.
+
+    >>> import tempfile
+    >>> d = tempfile.mkdtemp()
+    >>> atomic_write(os.path.join(d, "a"), [b"x", b"y"], perms=0o600, fsync=True)
+    >>> open(os.path.join(d, "a"), "rb").read()
+    b'xy'
+    >>> def boom():
+    ...     yield b"partial"
+    ...     raise RuntimeError("producer died")
+    >>> atomic_write(os.path.join(d, "a"), boom())
+    Traceback (most recent call last):
+      ...
+    RuntimeError: producer died
+    >>> open(os.path.join(d, "a"), "rb").read(), os.listdir(d)  # old content, no litter
+    (b'xy', ['a'])
+    """
+    if isinstance(data, str):
+        raise TypeError(
+            "atomic_write takes bytes (or chunks of bytes); encode text first"
+        )
+    path = os.path.realpath(os.fspath(path))
+    chunks = [data] if isinstance(data, _BYTES_LIKE) else data
+    dirname, name = os.path.split(path)
+    tmp = os.path.join(dirname, f".{name[:100]}.{uuid.uuid4().hex[:8]}.tmp")
+    mode = _mode_to_apply(path, perms)
+    try:
+        fd = os.open(
+            tmp, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o666 if mode is None else mode
+        )
+        with os.fdopen(fd, "wb") as fh:
+            if mode is not None and hasattr(os, "fchmod"):
+                os.fchmod(fh.fileno(), mode)  # exact, not narrowed by the umask
+            for chunk in chunks:
+                fh.write(chunk)
+            if fsync:
+                fh.flush()
+                os.fsync(fh.fileno())
+        os.replace(tmp, path)
+    except BaseException:
+        try:
+            os.unlink(tmp)
+        except OSError:
+            pass
+        raise
+    if fsync:
+        _fsync_dir(dirname)

dol/tests/test_atomic.py

@@ -0,0 +1,129 @@
+"""Tests for ``dol.atomic.atomic_write``."""
+
+import os
+import stat
+import sys
+
+import pytest
+
+from dol import atomic_write
+
+posix_only = pytest.mark.skipif(
+    sys.platform == "win32", reason="POSIX mode bits/symlinks"
+)
+
+
+def _mode(p):
+    return stat.S_IMODE(os.stat(p).st_mode)
+
+
+def test_bytes_and_chunks_and_overwrite(tmp_path):
+    p = tmp_path / "f"
+    atomic_write(p, b"one")
+    atomic_write(p, iter([b"t", bytearray(b"w"), memoryview(b"o")]))
+    assert p.read_bytes() == b"two"
+    assert os.listdir(tmp_path) == ["f"]
+
+
+def test_empty_data(tmp_path):
+    atomic_write(tmp_path / "f", iter([]))
+    assert (tmp_path / "f").read_bytes() == b""
+
+
+def test_str_is_refused_clearly(tmp_path):
+    with pytest.raises(TypeError, match="encode"):
+        atomic_write(tmp_path / "f", "text")
+    assert os.listdir(tmp_path) == []
+
+
+def test_failure_keeps_old_content_and_cleans_tmp(tmp_path):
+    p = tmp_path / "f"
+    atomic_write(p, b"old")
+
+    def gen():
+        yield b"new"
+        raise ValueError
+
+    with pytest.raises(ValueError):
+        atomic_write(p, gen())
+    assert p.read_bytes() == b"old" and os.listdir(tmp_path) == ["f"]
+
+
+def test_missing_parent_raises_and_leaves_nothing(tmp_path):
+    with pytest.raises(OSError):
+        atomic_write(tmp_path / "nope" / "f", b"x")
+    assert os.listdir(tmp_path) == []
+
+
+def test_tmp_is_hidden_unique_and_in_same_folder(tmp_path, monkeypatch):
+    seen = []
+    real = os.replace
+    monkeypatch.setattr(os, "replace", lambda a, b: (seen.append(a), real(a, b)))
+    atomic_write(tmp_path / "f", b"x")
+    atomic_write(tmp_path / "f", b"x")
+    assert len(set(seen)) == 2
+    for s in seen:
+        assert os.path.dirname(s) == os.path.realpath(tmp_path)
+        assert os.path.basename(s).startswith(".f.")
+
+
+def test_long_names_do_not_overflow(tmp_path):
+    name = "n" * 240
+    atomic_write(tmp_path / name, b"x")
+    assert (tmp_path / name).read_bytes() == b"x"
+
+
+def test_fsync_only_when_asked(tmp_path, monkeypatch):
+    calls = []
+    real = os.fsync
+    monkeypatch.setattr(os, "fsync", lambda fd: (calls.append(fd), real(fd)))
+    atomic_write(tmp_path / "f", b"x")
+    assert calls == []
+    atomic_write(tmp_path / "f", b"x", fsync=True)
+    assert len(calls) >= 1 and (tmp_path / "f").read_bytes() == b"x"
+
+
+@posix_only
+def test_perms_exact_despite_umask(tmp_path):
+    old = os.umask(0o077)
+    try:
+        atomic_write(tmp_path / "f", b"x", perms=0o644)
+    finally:
+        os.umask(old)
+    assert _mode(tmp_path / "f") == 0o644
+
+
+@posix_only
+def test_perms_set_before_data_is_written(tmp_path, monkeypatch):
+    monkeypatch.setattr(os, "fchmod", lambda *a: None)  # creation alone must do it
+    seen = []
+
+    def gen():
+        (tmp,) = [p for p in os.listdir(tmp_path) if p.startswith(".f.")]
+        seen.append(_mode(tmp_path / tmp))
+        yield b"secret"
+
+    old = os.umask(0)
+    try:
+        atomic_write(tmp_path / "f", gen(), perms=0o600)
+    finally:
+        os.umask(old)
+    assert seen == [0o600]
+
+
+@posix_only
+def test_existing_mode_preserved(tmp_path):
+    p = tmp_path / "f"
+    p.write_bytes(b"old")
+    os.chmod(p, 0o600)
+    atomic_write(p, b"new")
+    assert _mode(p) == 0o600 and p.read_bytes() == b"new"
+
+
+@posix_only
+def test_symlink_is_followed(tmp_path):
+    target, link = tmp_path / "target", tmp_path / "link"
+    target.write_bytes(b"old")
+    link.symlink_to(target)
+    atomic_write(link, b"new")
+    assert link.is_symlink() and target.read_bytes() == b"new"