Python - Atomic file Reader/Writer

atomic_file.py

import os
import tempfile
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, IO, Type


class AtomicFile:
    @dataclass(eq=False, kw_only=True)
    class Reader:
        """
        AtomicFile provides a mechanism for atomically reading a file.

        This class acts as a context manager, ensuring that a file's parent
        directory is flushed before giving access to the data.

        Typical usage:

        ```
        with AtomicFile.Reader('filename.txt') as f:
            foo = f.read()
        ```

        Parameters:
        - filename (str): Path to the target file that needs atomic replacement.
        """
        filename: Path
        file: IO[str] | None = field(init=False, default=None)

        def __enter__(self) -> IO[str]:
            # Ensure output directory is flushed
            fd = os.open(self.filename.parent, os.O_CLOEXEC | os.O_DIRECTORY | os.O_RDONLY)

            try:
                os.fsync(fd)
            finally:
                os.close(fd)

            self.file = open(self.filename, 'r', encoding='utf-8')

            return self.file

        def __exit__(self, exc_type: Type[BaseException], *_: Any) -> None:
            assert self.file
            self.file.close()

    @dataclass(eq=False, kw_only=True)
    class Writer:
        """
        AtomicFile provides a mechanism for atomically writing a file.

        This class acts as a context manager, ensuring that a file's content
        is either fully replaced or, in the case of errors, left unchanged.
        The atomic replacement is achieved by first writing to a temporary file
        and then renaming the temporary file to replace the original file.

        Typical usage:

        ```
        with AtomicFile.Writer('filename.txt') as f:
            print('new content', file=f)
        ```

        Parameters:
        - filename (str): Path to the target file that needs atomic replacement.
        - mode (str, optional): Mode in which the file is opened, default is 'w'.

        Note:
        - The temporary file is created in the same directory as the original file
          to ensure atomicity of `os.rename()`.
        """
        filename: Path
        mode: str = 'w'
        temp_file: IO[str] | None = field(default=None, init=False)

        def __enter__(self) -> IO[str]:
            # Create a temporary file in the same directory as the target file
            self.temp_file = tempfile.NamedTemporaryFile(
                prefix='.tmp.',
                mode=self.mode,
                dir=self.filename.parent,
                delete=False,
                encoding='utf-8',
            )

            # If the target file already exists, copy its mode (permissions)
            try:
                mode = self.filename.stat().st_mode
            except FileNotFoundError:
                mode = 0o644

            os.chmod(self.temp_file.name, mode)

            return self.temp_file

        def __exit__(self, exc_type: Type[BaseException], *_: Any) -> None:
            assert self.temp_file is not None

            if exc_type:
                # Clean up temp file in case of an error
                self.temp_file.close()
                os.remove(self.temp_file.name)
                return

            # Only finalize if there was no exception inside the with-block
            self.temp_file.flush()
            os.fsync(self.temp_file.fileno())
            self.temp_file.close()

            # Atomically rename the temporary file to replace the original file
            os.rename(self.temp_file.name, self.filename)