Write user files and media atomically

src/backend/utils.py

@@ -8,16 +8,23 @@ import inspect
 import io
 import json
 import xml.etree.cElementTree as xml_etree  # FIXME: bandit warning
+from contextlib import asynccontextmanager
 from datetime import datetime, timedelta
 from enum import Enum
 from enum import auto as autostr
 from pathlib import Path
+from tempfile import NamedTemporaryFile
 from types import ModuleType
-from typing import Any, Dict, Mapping, Sequence, Tuple, Type
+from typing import (
+    Any, AsyncIterator, Dict, Mapping, Sequence, Tuple, Type, Union,
+)
 from uuid import UUID
 
+import aiofiles
 import filetype
 from aiofiles.threadpool.binary import AsyncBufferedReader
+from aiofiles.threadpool.text import AsyncTextIOWrapper
+
 from nio.crypto import AsyncDataT as File
 from nio.crypto import async_generator_from_data
 
@@ -185,3 +192,21 @@ def classes_defined_in(module: ModuleType) -> Dict[str, Type]:
         if not m[0].startswith("_") and
         m[1].__module__.startswith(module.__name__)
     }
+
+
+@asynccontextmanager
+async def atomic_write(
+    path: Union[Path, str], binary: bool = False, **kwargs,
+) -> AsyncIterator[Union[AsyncTextIOWrapper, AsyncBufferedReader]]:
+    """Write a file asynchronously (using aiofiles) and atomically."""
+
+    mode      = "wb" if binary else "w"
+    path      = Path(path)
+    temp      = NamedTemporaryFile(dir=path.parent, delete=False)
+    temp_path = Path(temp.name)
+
+    try:
+        async with aiofiles.open(temp_path, mode, **kwargs) as out:
+            yield out
+    finally:
+        temp_path.replace(path)