beam: atomic_write module + tests (crash-safe fsync+rename)

wizard-work/beam/atomic_write.py

@@ -0,0 +1,92 @@
+"""Atomic, crash-safe file writing.
+
+Writes to a temporary file in the *same* directory as the target (so the
+rename is atomic on the same filesystem), fsyncs both the file and the
+directory, and atomically renames into place. On any exception the temp
+file is removed and the original (if any) is left untouched.
+
+Usage:
+    from atomic_write import atomic_write
+    with atomic_write("out.txt") as f:
+        f.write("hello")
+    # file is durable + atomic on return
+
+Also supports mode="a" (append) by reading existing content? No — append is
+NOT atomic in the same sense; use mode="w" for full replace. We keep a
+single mode for correctness.
+"""
+from __future__ import annotations
+
+import os
+import tempfile
+from contextlib import contextmanager
+from typing import IO, Iterator, Optional
+
+
+@contextmanager
+def atomic_write(
+    path: str,
+    *,
+    mode: str = "w",
+    encoding: Optional[str] = "utf-8",
+    fsync: bool = True,
+    preserve_mode: bool = True,
+) -> Iterator[IO]:
+    """Context manager yielding a writable file object; commit on clean exit."""
+    if mode not in ("w", "wb"):
+        raise ValueError(f"unsupported mode {mode!r}; use 'w' or 'wb'")
+
+    directory = os.path.dirname(os.path.abspath(path)) or "."
+    os.makedirs(directory, exist_ok=True)
+
+    # Preserve existing file permissions when overwriting.
+    existing_mode = None
+    if preserve_mode and os.path.exists(path):
+        existing_mode = os.stat(path).st_mode & 0o777
+
+    fd, tmp = tempfile.mkstemp(prefix=".atomic_", dir=directory)
+    try:
+        # Apply mode bits up front.
+        if existing_mode is not None:
+            os.fchmod(fd, existing_mode)
+        f = os.fdopen(fd, mode, encoding=encoding if "b" not in mode else None, closefd=True)
+        # We own f now; fd is closed by fdopen.
+        fd = -1
+        try:
+            yield f
+            f.flush()
+            if fsync:
+                os.fsync(f.fileno())
+            f.close()
+            # Atomic replace.
+            os.replace(tmp, path)
+            tmp = None  # committed
+            if fsync:
+                # fsync directory so the rename is durable.
+                dirfd = os.open(directory, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0))
+                try:
+                    os.fsync(dirfd)
+                finally:
+                    os.close(dirfd)
+        except BaseException:
+            try:
+                f.close()
+            except Exception:
+                pass
+            raise
+    finally:
+        if tmp is not None and os.path.exists(tmp):
+            try:
+                os.unlink(tmp)
+            except OSError:
+                pass
+
+
+def atomic_write_text(path: str, data: str, *, encoding: str = "utf-8", fsync: bool = True) -> None:
+    with atomic_write(path, encoding=encoding, fsync=fsync) as f:
+        f.write(data)
+
+
+def atomic_write_bytes(path: str, data: bytes, *, fsync: bool = True) -> None:
+    with atomic_write(path, mode="wb", encoding=None, fsync=fsync) as f:
+        f.write(data)

wizard-work/beam/test_atomic_write.py

@@ -0,0 +1,86 @@
+import os
+import stat
+import tempfile
+import unittest
+
+from atomic_write import atomic_write, atomic_write_text, atomic_write_bytes
+
+
+class TestAtomicWrite(unittest.TestCase):
+    def setUp(self):
+        self.d = tempfile.mkdtemp()
+        self.p = os.path.join(self.d, "out.txt")
+
+    def test_basic_text(self):
+        with atomic_write(self.p) as f:
+            f.write("hello world")
+        with open(self.p) as f:
+            self.assertEqual(f.read(), "hello world")
+
+    def test_bytes(self):
+        atomic_write_bytes(self.p, b"\x00\x01\x02")
+        with open(self.p, "rb") as f:
+            self.assertEqual(f.read(), b"\x00\x01\x02")
+
+    def test_text_helper(self):
+        atomic_write_text(self.p, "abc")
+        with open(self.p) as f:
+            self.assertEqual(f.read(), "abc")
+
+    def test_no_partial_on_exception(self):
+        # pre-existing content must survive a failed write
+        with open(self.p, "w") as f:
+            f.write("ORIGINAL")
+        try:
+            with atomic_write(self.p) as f:
+                f.write("PARTIAL")
+                raise RuntimeError("boom")
+        except RuntimeError:
+            pass
+        with open(self.p) as f:
+            self.assertEqual(f.read(), "ORIGINAL")
+
+    def test_no_temp_leftover_on_exception(self):
+        try:
+            with atomic_write(self.p) as f:
+                f.write("x")
+                raise ValueError("x")
+        except ValueError:
+            pass
+        leftovers = [n for n in os.listdir(self.d) if n.startswith(".atomic_")]
+        self.assertEqual(leftovers, [])
+
+    def test_preserves_mode(self):
+        with open(self.p, "w") as f:
+            f.write("x")
+        os.chmod(self.p, 0o600)
+        with atomic_write(self.p) as f:
+            f.write("y")
+        self.assertEqual(stat.S_IMODE(os.stat(self.p).st_mode), 0o600)
+
+    def test_overwrite_replaces(self):
+        atomic_write_text(self.p, "first")
+        atomic_write_text(self.p, "second")
+        with open(self.p) as f:
+            self.assertEqual(f.read(), "second")
+
+    def test_creates_parent_dir(self):
+        p = os.path.join(self.d, "sub", "deep", "f.txt")
+        atomic_write_text(p, "nested")
+        with open(p) as f:
+            self.assertEqual(f.read(), "nested")
+
+    def test_rejects_bad_mode(self):
+        with self.assertRaises(ValueError):
+            with atomic_write(self.p, mode="a"):
+                pass
+
+    def test_empty_write(self):
+        atomic_write_text(self.p, "")
+        self.assertTrue(os.path.exists(self.p))
+        with open(self.p) as f:
+            self.assertEqual(f.read(), "")
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)