diff --git a/src/tagstudio/previews/renderers/archive.py b/src/tagstudio/previews/renderers/archive.py index 1f1b612cf..727c27ab1 100644 --- a/src/tagstudio/previews/renderers/archive.py +++ b/src/tagstudio/previews/renderers/archive.py @@ -2,17 +2,17 @@ # SPDX-License-Identifier: GPL-3.0-only -import tarfile -import zipfile from io import BytesIO from pathlib import Path -from typing import Literal, override +from tarfile import TarFile +from typing import Self, override +from zipfile import ZipFile -import py7zr -import py7zr.io -import rarfile import structlog from PIL.Image import Image +from py7zr import SevenZipFile +from py7zr.io import BytesIOFactory +from rarfile import RarFile from tagstudio.core.enums import Theme from tagstudio.core.media_types import MediaTypes @@ -22,7 +22,7 @@ logger = structlog.get_logger(__name__) -type Archive = zipfile.ZipFile | rarfile.RarFile | SevenZipFile | TarFile +type Archive = Rar | SevenZip | Tar | Zip class ArchivePreview(BasePreview): @@ -57,41 +57,50 @@ def render( return archive_thumb(filepath) -class SevenZipFile(py7zr.SevenZipFile): - """Wrapper around py7zr.SevenZipFile to mimic zipfile.ZipFile's API.""" +class Rar(RarFile): + """Wrapper around RarFile for a unified API.""" - def __init__(self, filepath: Path, mode: Literal["r"]) -> None: - super().__init__(filepath, mode) + @classmethod + def open_archive(cls, filepath: Path) -> Self: + return cls(filepath, "r") + + +class SevenZip(SevenZipFile): + """Wrapper around SevenZipFile for a unified API.""" + + @classmethod + def open_archive(cls, filepath: Path) -> Self: + return cls(filepath, "r") def read(self, name: str) -> bytes: # SevenZipFile must be reset after every extraction # See https://py7zr.readthedocs.io/en/stable/api.html#py7zr.SevenZipFile.extract self.reset() - factory = py7zr.io.BytesIOFactory(limit=10485760) # 10 MiB + factory = BytesIOFactory(limit=10485760) # 10 MiB self.extract(targets=[name], factory=factory) return factory.get(name).read() -class TarFile: - """Wrapper around tarfile.TarFile to mimic zipfile.ZipFile's API.""" +class Tar(TarFile): + """Wrapper around TarFile for a unified API.""" - def __init__(self, filepath: Path, mode: Literal["r"]) -> None: - self.tar: tarfile.TarFile - self.filepath = filepath - self.mode: Literal["r"] = mode + @classmethod + def open_archive(cls, filepath: Path) -> Self: + return cls.open(filepath, "r") def namelist(self) -> list[str]: - return self.tar.getnames() + return self.getnames() def read(self, name: str) -> bytes: - return unwrap(self.tar.extractfile(name)).read() + return unwrap(self.extractfile(name)).read() - def __enter__(self) -> TarFile: - self.tar = tarfile.open(name=self.filepath, mode=self.mode).__enter__() - return self - def __exit__(self, *args) -> None: # pyright: ignore[reportUnknownParameterType, reportMissingParameterType] - self.tar.__exit__(*args) +class Zip(ZipFile): + """Wrapper around ZipFile for a unified API.""" + + @classmethod + def open_archive(cls, filepath: Path) -> Self: + return cls(filepath, "r") def open_archive(filepath: Path) -> Archive: @@ -104,11 +113,11 @@ def open_archive(filepath: Path) -> Archive: Archive: The opened archive. """ exts = MediaTypes.get_equivalent_exts(MediaTypes.get_ext(filepath)) - archiver: type[Archive] = zipfile.ZipFile + archiver: type[Archive] = Zip if exts & {".7z", ".cb7", ".s7z"}: - archiver = SevenZipFile + archiver = SevenZip elif exts & {".cbr", ".rar"}: - archiver = rarfile.RarFile + archiver = Rar elif exts & { ".cbt", ".tar", @@ -119,8 +128,8 @@ def open_archive(filepath: Path) -> Archive: ".tar.xz", ".tar.zst", }: - archiver = TarFile - return archiver(filepath, "r") + archiver = Tar + return archiver.open_archive(filepath) def first_image_in_archive(archive: Archive) -> Image | None: @@ -151,7 +160,6 @@ def archive_thumb( Args: filepath (Path): The path to the archive. image_names: (list[Path] | list[str] | None): List of embedded image names to search for. - ext (str): The file extension. Used to help determine more specific archive type. Returns: Image: The first image found in the archive.