feat: Expand tarball thumbnail support

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.