260705_2

common_util/common_util_atomic.py

@@ -0,0 +1,55 @@
+"""바이너리·npz 파일을 원자적으로 저장하는 공통 유틸리티."""
+
+import os
+import tempfile
+from pathlib import Path
+
+import numpy as np
+
+
+def atomic_write_bytes(path: str | Path, payload: bytes) -> None:
+    """같은 디렉터리의 임시 파일을 교체하여 바이트를 원자적으로 저장한다."""
+    target = Path(path)
+    target.parent.mkdir(parents=True, exist_ok=True)
+    temporary_path: Path | None = None
+    try:
+        with tempfile.NamedTemporaryFile(
+            mode="wb",
+            dir=target.parent,
+            prefix=f".{target.name}.",
+            suffix=".tmp",
+            delete=False,
+        ) as temporary:
+            temporary.write(payload)
+            temporary.flush()
+            os.fsync(temporary.fileno())
+            temporary_path = Path(temporary.name)
+        os.replace(temporary_path, target)
+        temporary_path = None
+    finally:
+        if temporary_path is not None:
+            temporary_path.unlink(missing_ok=True)
+
+
+def atomic_write_npz(path: str | Path, **arrays: np.ndarray) -> None:
+    """같은 디렉터리의 임시 파일을 교체하여 압축 npz를 원자적으로 저장한다."""
+    target = Path(path)
+    target.parent.mkdir(parents=True, exist_ok=True)
+    temporary_path: Path | None = None
+    try:
+        with tempfile.NamedTemporaryFile(
+            mode="wb",
+            dir=target.parent,
+            prefix=f".{target.name}.",
+            suffix=".tmp",
+            delete=False,
+        ) as temporary:
+            np.savez_compressed(temporary, **arrays)
+            temporary.flush()
+            os.fsync(temporary.fileno())
+            temporary_path = Path(temporary.name)
+        os.replace(temporary_path, target)
+        temporary_path = None
+    finally:
+        if temporary_path is not None:
+            temporary_path.unlink(missing_ok=True)