93 lines
2.9 KiB
Python
93 lines
2.9 KiB
Python
"""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)
|