mapdet3d/data/io/tar.py

"""Tar data backend.

This backend works with filepaths pointing to files inside tar archives.
Similar to HDF5Backend but for tar files.
"""

from __future__ import annotations

import os
import tarfile
from typing import Literal

from .base import DataBackend


class TarBackend(DataBackend):
    """Backend for loading data from tar files.

    This backend works with filepaths pointing to files inside tar archives.
    The filepath format should be: 'path/to/file.tar/internal/path/to/file'

    Examples:
        >>> backend = TarBackend()
        >>> data = backend.get("dataset.tar/images/img001.png")
        >>> files = backend.listdir("dataset.tar/images")
    """

    def __init__(self) -> None:
        """Creates an instance of the class."""
        super().__init__()
        self.tar_cache: dict[str, tarfile.TarFile] = {}

    @staticmethod
    def _get_tar_path(
        filepath: str, allow_omitted_ext: bool = True
    ) -> tuple[str, str]:
        """Get .tar path and internal path from filepath.

        Args:
            filepath (str): The filepath to retrieve the data from.
                Should have format: 'path/to/file.tar/internal/path'
            allow_omitted_ext (bool, optional): Whether to allow omitted
                extension, in which case the backend will try to append
                '.tar' to the filepath. Defaults to True.

        Returns:
            tuple[str, str]: The .tar path and the internal path.

        Examples:
            >>> TarBackend._get_tar_path("path/to/file.tar/key1/key2")
            ("path/to/file.tar", "key1/key2")
            >>> TarBackend._get_tar_path("path/to/file/key1/key2", True)
            ("path/to/file.tar", "key1/key2")  # if file.tar exists
        """
        filepath_as_list = filepath.split("/")
        internal_parts = []

        tar_path = filepath
        while True:
            if tar_path.endswith(".tar") or tar_path == "":
                break
            if allow_omitted_ext and os.path.exists(tar_path + ".tar"):
                tar_path = tar_path + ".tar"
                break
            if internal_parts or not os.path.exists(tar_path):
                internal_parts.insert(0, filepath_as_list.pop())
                tar_path = "/".join(filepath_as_list)
            else:
                break

        internal_path = "/".join(internal_parts)
        return tar_path, internal_path

    def exists(self, filepath: str) -> bool:
        """Check if filepath exists.

        Args:
            filepath (str): Path to file inside tar.

        Returns:
            bool: True if file exists, False otherwise.
        """
        tar_path, internal_path = self._get_tar_path(filepath)
        if not os.path.exists(tar_path):
            return False

        try:
            tar_file = self._get_client(tar_path)
            try:
                tar_file.getmember(internal_path)
                return True
            except KeyError:
                return False
        except (tarfile.TarError, FileNotFoundError):
            return False

    def _get_client(self, tar_path: str) -> tarfile.TarFile:
        """Get TarFile client from path.

        Args:
            tar_path (str): Path to tar file.

        Returns:
            tarfile.TarFile: The opened tar file.
        """
        if tar_path not in self.tar_cache:
            client = tarfile.open(tar_path, mode="r:*")
            self.tar_cache[tar_path] = client
        return self.tar_cache[tar_path]

    def get(self, filepath: str) -> bytes:
        """Get file content as bytes.

        Args:
            filepath (str): The path to the file. It consists of a tar path
                together with the internal path, e.g.: "/path/to/file.tar/
                internal/path/data.json".

        Raises:
            FileNotFoundError: If tar file doesn't exist.
            ValueError: If internal path not found inside tar.

        Returns:
            bytes: The file content in bytes
        """
        tar_path, internal_path = self._get_tar_path(filepath)

        if not os.path.exists(tar_path):
            raise FileNotFoundError(f"Tar file not found: {tar_path}")

        tar_file = self._get_client(tar_path)

        try:
            member = tar_file.getmember(internal_path)
            if member.isdir():
                raise ValueError(f"{internal_path} is a directory, not a file")

            file_obj = tar_file.extractfile(member)
            if file_obj is None:
                raise ValueError(f"Cannot extract file: {internal_path}")

            return file_obj.read()

        except KeyError:
            raise ValueError(f"Path {internal_path} not found in {tar_path}")

    def isfile(self, filepath: str) -> bool:
        """Check if filepath is a file (not a directory).

        Args:
            filepath (str): Path to file inside tar.

        Raises:
            FileNotFoundError: If tar file doesn't exist.
            ValueError: If path not found inside tar.

        Returns:
            bool: True if it's a file, False if it's a directory.
        """
        tar_path, internal_path = self._get_tar_path(filepath)

        if not os.path.exists(tar_path):
            raise FileNotFoundError(f"Tar file not found: {tar_path}")

        tar_file = self._get_client(tar_path)

        try:
            member = tar_file.getmember(internal_path)
            return member.isfile()
        except KeyError:
            raise ValueError(f"Path {internal_path} not found in {tar_path}")

    def listdir(self, filepath: str) -> list[str]:
        """List all files and directories in the given path.

        Args:
            filepath (str): Path to directory inside tar.

        Raises:
            FileNotFoundError: If tar file doesn't exist.
            ValueError: If path not found or is not a directory.

        Returns:
            list[str]: List of file/directory names (not full paths).
        """
        tar_path, internal_path = self._get_tar_path(filepath)

        if not os.path.exists(tar_path):
            raise FileNotFoundError(f"Tar file not found: {tar_path}")

        tar_file = self._get_client(tar_path)

        # Normalize path
        internal_path = internal_path.rstrip("/")
        if internal_path and not internal_path.endswith("/"):
            prefix = internal_path + "/"
        else:
            prefix = internal_path

        # Collect immediate children
        children = set()
        for member in tar_file.getmembers():
            if member.name.startswith(prefix):
                # Get relative path
                relative = member.name[len(prefix) :]

                # Only get immediate children
                if relative and "/" not in relative.rstrip("/"):
                    children.add(relative.rstrip("/"))
                elif "/" in relative:
                    # Add the directory name only
                    dir_name = relative.split("/")[0]
                    children.add(dir_name)

        return sorted(children)

    def set(
        self, filepath: str, content: bytes, mode: Literal["w", "a"] = "a"
    ) -> None:
        """Set file content (not implemented for tar files).

        Tar files are typically read-only. Use tar command line tools
        to create/modify tar files.

        Raises:
            NotImplementedError: Always raised.
        """
        raise NotImplementedError(
            "Writing to tar files is not supported. "
            "Use tar command line tools to create tar archives."
        )

    def close(self) -> None:
        """Close all opened tar files."""
        for tar_file in self.tar_cache.values():
            tar_file.close()
        self.tar_cache.clear()

    def __del__(self):
        """Cleanup when object is destroyed."""
        self.close()