bookwyrm/bookwyrm/utils/tar.py

"""manage tar files for user exports"""
import io
import os
import tarfile
from typing import Any, Optional
from uuid import uuid4
from django.core.files import File


class BookwyrmTarFile(tarfile.TarFile):
    """Create tar files for user exports"""

    def write_bytes(self, data: bytes) -> None:
        """Add a file containing bytes to the archive"""
        buffer = io.BytesIO(data)
        info = tarfile.TarInfo("archive.json")
        info.size = len(data)
        self.addfile(info, fileobj=buffer)

    def add_image(
        self, image: Any, filename: Optional[str] = None, directory: str = ""
    ) -> None:
        """
        Add an image to the tar archive
        :param str filename: overrides the file name set by image
        :param str directory: the directory in the archive to put the image
        """
        if filename is None:
            dst_filename = image.name
        else:
            dst_filename = filename + os.path.splitext(image.name)[1]
        dst_path = os.path.join(directory, dst_filename)

        info = tarfile.TarInfo(name=dst_path)
        info.size = image.size

        self.addfile(info, fileobj=image)

    def read(self, filename: str) -> Any:
        """read data from the tar"""
        if reader := self.extractfile(filename):
            return reader.read()
        return None

    def write_image_to_file(self, filename: str, file_field: Any) -> None:
        """add an image to the tar"""
        extension = os.path.splitext(filename)[1]
        if buf := self.extractfile(filename):
            filename = str(uuid4()) + extension
            file_field.save(filename, File(buf))