diff --git a/wizard-work/beam/atomic_write.py b/wizard-work/beam/atomic_write.py new file mode 100644 index 0000000..9b888f8 --- /dev/null +++ b/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) diff --git a/wizard-work/beam/test_atomic_write.py b/wizard-work/beam/test_atomic_write.py new file mode 100644 index 0000000..703b346 --- /dev/null +++ b/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)