crawler-intra-mart/storage/file_writer.py

import json
import os
import shutil
import tempfile

from pathlib import Path

import aiofiles

from utils.logger import logger


class FileWriter:
    """
    Production-ready file writer.

    Supports:
    - text
    - json
    - binary
    - async write
    - atomic write
    - safe overwrite
    """

    def __init__(
        self,
        encoding="utf-8"
    ):
        self.encoding = encoding

    def ensure_parent(
        self,
        path: Path
    ):
        """
        Ensure parent directory exists.
        """

        path.parent.mkdir(
            parents=True,
            exist_ok=True
        )

    def write_text(
        self,
        path: Path,
        content: str,
        overwrite=True
    ):
        """
        Write text file safely.
        """

        self.ensure_parent(path)

        if path.exists() and not overwrite:

            logger.warning(
                f"Skip existing file: {path}"
            )

            return

        # Create the temp file in the SAME directory as the
        # destination so that shutil.move can do a same-FS
        # atomic rename. mkstemp on /tmp can land on a
        # different mount point, forcing a non-atomic
        # copy+delete fallback. Also: mkstemp returns a raw
        # fd that must be closed -- the old code leaked one
        # fd per write_text call.
        temp_fd, temp_path = tempfile.mkstemp(
            prefix=".tmp_",
            suffix=path.suffix,
            dir=str(path.parent)
        )

        try:

            with os.fdopen(
                temp_fd,
                "w",
                encoding=self.encoding
            ) as f:

                f.write(content)

            shutil.move(
                temp_path,
                path
            )

            logger.info(
                f"Saved text file: {path}"
            )

        except Exception as e:

            logger.exception(
                f"Failed writing text "
                f"{path}: {e}"
            )

            # Clean up the temp file if move never happened.
            try:
                Path(temp_path).unlink(missing_ok=True)
            except Exception:
                pass

            raise

    async def write_text_async(
        self,
        path: Path,
        content: str
    ):
        """
        Async text writer.
        """

        self.ensure_parent(path)

        try:

            async with aiofiles.open(
                path,
                "w",
                encoding=self.encoding
            ) as f:

                await f.write(content)

            logger.info(
                f"Saved async text: {path}"
            )

        except Exception as e:

            logger.exception(
                f"Async write failed "
                f"{path}: {e}"
            )

            raise

    def write_json(
        self,
        path: Path,
        data,
        indent=2
    ):
        """
        Write JSON file.
        """

        self.ensure_parent(path)

        try:

            with open(
                path,
                "w",
                encoding=self.encoding
            ) as f:

                json.dump(
                    data,
                    f,
                    ensure_ascii=False,
                    indent=indent
                )

            logger.info(
                f"Saved JSON: {path}"
            )

        except Exception as e:

            logger.exception(
                f"JSON write failed "
                f"{path}: {e}"
            )

            raise

    async def write_json_async(
        self,
        path: Path,
        data,
        indent=2
    ):
        """
        Async JSON writer.
        """

        self.ensure_parent(path)

        try:

            content = json.dumps(
                data,
                ensure_ascii=False,
                indent=indent
            )

            async with aiofiles.open(
                path,
                "w",
                encoding=self.encoding
            ) as f:

                await f.write(content)

            logger.info(
                f"Saved async JSON: {path}"
            )

        except Exception as e:

            logger.exception(
                f"Async JSON write failed "
                f"{path}: {e}"
            )

            raise

    def write_binary(
        self,
        path: Path,
        content: bytes
    ):
        """
        Write binary file.
        """

        self.ensure_parent(path)

        try:

            with open(
                path,
                "wb"
            ) as f:

                f.write(content)

            logger.info(
                f"Saved binary: {path}"
            )

        except Exception as e:

            logger.exception(
                f"Binary write failed "
                f"{path}: {e}"
            )

            raise

    async def append_text(
        self,
        path: Path,
        content: str
    ):
        """
        Append text to file.
        """

        self.ensure_parent(path)

        try:

            async with aiofiles.open(
                path,
                "a",
                encoding=self.encoding
            ) as f:

                await f.write(content)

            logger.info(
                f"Appended text: {path}"
            )

        except Exception as e:

            logger.exception(
                f"Append failed "
                f"{path}: {e}"
            )

            raise

    def read_text(
        self,
        path: Path
    ):
        """
        Read text file.
        """

        try:

            with open(
                path,
                "r",
                encoding=self.encoding
            ) as f:

                return f.read()

        except Exception as e:

            logger.exception(
                f"Read failed "
                f"{path}: {e}"
            )

            raise

    def exists(
        self,
        path: Path
    ):
        """
        Check file existence.
        """

        return path.exists()

    def delete(
        self,
        path: Path
    ):
        """
        Delete file safely.
        """

        try:

            if path.exists():

                path.unlink()

                logger.info(
                    f"Deleted file: {path}"
                )

        except Exception as e:

            logger.exception(
                f"Delete failed "
                f"{path}: {e}"
            )

    def sanitize_filename(
        self,
        filename: str
    ):
        """
        Remove invalid filesystem chars.
        """

        invalid_chars = [
            "<",
            ">",
            ":",
            "\"",
            "/",
            "\\",
            "|",
            "?",
            "*"
        ]

        for char in invalid_chars:

            filename = filename.replace(
                char,
                "_"
            )

        return filename.strip()

    def write_markdown_with_metadata(
        self,
        path: Path,
        markdown: str,
        metadata: dict = None
    ):
        """
        Write markdown with YAML frontmatter.
        """

        self.ensure_parent(path)

        content = ""

        if metadata:

            content += "---\n"

            for key, value in metadata.items():

                content += f"{key}: {value}\n"

            content += "---\n\n"

        content += markdown

        self.write_text(
            path,
            content
        )