fix: proper compound ext support + compression ext handling (#1536)

* fix: properly handle compound file extentions

* fix(ui): register other supported tar types to render previews

Co-authored-by: Sola-ris <190788035+Sola-ris@users.noreply.github.com>

* chore: correct stale docstring

---------

Co-authored-by: Sola-ris <190788035+Sola-ris@users.noreply.github.com>

Author: CyanVoxel

docs/preview-support.md

@@ -105,12 +105,17 @@ Preview support for office documents or well-known project file formats varies b
 
 Archive thumbnails will display the first image from the archive within the Preview Panel.
 
-| Filetype | Extensions     |
-| -------- | -------------- |
-| 7-Zip    | `.7z`, `.s7z`  |
-| RAR      | `.rar`         |
-| Tar      | `.tar`, `.tgz` |
-| Zip      | `.zip`         |
+| Filetype        | Extensions                                             |
+| --------------- | ------------------------------------------------------ |
+| 7-Zip           | `.7z`, `.s7z`                                          |
+| RAR             | `.rar`                                                 |
+| Tar             | `.tar`                                                 |
+| Tar (bzip2)     | `.tar.bz2`, `.tar.bz`, `.tbz2`, `.tbz`, `.tb2`, `.tz2` |
+| Tar (gzip)      | `.tar.gz`, `.tgz`, `.taz`                              |
+| Tar (LZMA)      | `.tar.lzma`, `.tlz`                                    |
+| Tar (XZ)        | `.tar.xz`, `.txz`                                      |
+| Tar (Zstandard) | `.tar.zst`, `.tzst`                                    |
+| Zip             | `.zip`                                                 |
 
 ### :material-book: eBooks
 

src/tagstudio/core/library/alchemy/visitors.py

@@ -6,7 +6,7 @@
 from typing import TYPE_CHECKING, override
 
 import structlog
-from sqlalchemy import ColumnElement, and_, distinct, false, func, or_, select
+from sqlalchemy import ColumnElement, and_, distinct, false, func, literal, or_, select
 from sqlalchemy.orm import Session
 from sqlalchemy.sql.operators import ilike_op
 
@@ -100,17 +100,16 @@ def visit_constraint(self, node: Constraint) -> ColumnElement[bool]:
             extensions: set[str] = (
                 media_type.context_sets.get(SEARCH, set()) if media_type else set()
             )
-            return Entry.suffix.in_(map(lambda x: x.replace(".", ""), extensions))
+            return self._entry_has_any_ext(extensions)
 
         elif node.type == ConstraintType.FileType:
-            # NOTE: Entries store their suffix without a leading dot, the MediaTypes system includes
-            # the leading dot (if any), and this search system should take in either variant.
-            return or_(
-                *[
-                    Entry.suffix.ilike(ft.removeprefix("."))
-                    for ft in MediaTypes.get_equivalent_exts(f".{node.value.removeprefix('.')}")
-                ]
-            )
+            # NOTE: The MediaTypes system includes the leading dot (if any), and this search system
+            # should take in either variant.
+            ext = f".{node.value.removeprefix('.').lower()}"
+            extensions = MediaTypes.get_equivalent_exts(ext).copy()
+            for compound_ext in MediaTypes.get_compound_exts(ext):
+                extensions |= MediaTypes.get_equivalent_exts(compound_ext)
+            return self._entry_has_any_ext(extensions)
         elif node.type == ConstraintType.Special:  # noqa: SIM102 unnecessary once there is a second special constraint
             if node.value.lower() == "untagged":
                 return ~Entry.id.in_(select(Entry.id).join(TagEntry))
@@ -126,6 +125,30 @@ def visit_property(self, node: Property) -> ColumnElement[bool]:
     def visit_not(self, node: Not) -> ColumnElement[bool]:
         return ~self.visit(node.child)
 
+    def _entry_has_any_ext(self, extensions: set[str]) -> ColumnElement[bool]:
+        """Returns a Boolean Expression that is true if the Entry has any of the extensions."""
+        # NOTE: Entries only store their last suffix (lowercase and without a leading dot), so
+        # compound extensions also match the end of the path.
+        suffixes: set[str] = set()
+        last_suffixes_by_prefix: dict[str, set[str]] = {}
+        for ext in extensions:
+            prefix, _, last_suffix = ext.rpartition(".")
+            if prefix:
+                last_suffixes_by_prefix.setdefault(prefix, set()).add(last_suffix)
+            else:
+                suffixes.add(last_suffix)
+
+        return or_(
+            Entry.suffix.in_(suffixes),
+            *[
+                and_(
+                    Entry.suffix.in_(last_suffixes),
+                    Entry.path.like(literal(f"%{prefix}.") + Entry.suffix),
+                )
+                for prefix, last_suffixes in last_suffixes_by_prefix.items()
+            ],
+        )
+
     def __get_tag_ids(self, tag_name: str, include_children: bool = True) -> list[int]:
         """Given a tag name find the ids of all tags that this name could refer to."""
         with Session(self.lib.engine) as session:

src/tagstudio/core/media_types.py

@@ -3,6 +3,7 @@
 
 
 import re
+from pathlib import PurePath
 from typing import Any
 
 import structlog
@@ -105,6 +106,7 @@ class MediaTypes(metaclass=SanitizedAttr):
     """A singleton class that manages registered media types and their relationships."""
 
     _chained_groups: dict[str, set[str]] = {}
+    _compound_exts_by_last_suffix: dict[str, list[str]] = {}
     _name_to_key_map: dict[str, str] = {}
     all_groups: list[MediaTypeGroup] = []
     equivalent_exts: dict[str, set[str]] = {}
@@ -275,6 +277,16 @@ def register(cls, group_key: str, ext: str | list[str], contexts: str | list[str
             for e in ext:
                 cls.equivalent_exts.setdefault(e, set(ext))
 
+        # Index compound extensions (e.g. ".tar.gz") by their last suffix (.gz -> .tar.gz)
+        for e in ext:
+            if e.count(".") > 1:
+                compound_exts = cls._compound_exts_by_last_suffix.setdefault(
+                    f".{e.rsplit('.', 1)[1]}", []
+                )
+                if e not in compound_exts:
+                    compound_exts.append(e)
+                    compound_exts.sort(key=len, reverse=True)
+
         # Create any chained groups from dot notations (e.g. "adobe.photoshop")
         name_parts = group_key.split(".")
         for i in range(1, len(name_parts)):
@@ -301,5 +313,33 @@ def get_equivalent_exts(cls, ext: str) -> set[str]:
         """
         return cls.equivalent_exts.get(ext, {ext})
 
+    @classmethod
+    def get_ext(cls, path: PurePath) -> str:
+        """Return the longest registered extension a path ends with, or else its last suffix.
+
+        For example, `my.archive.tar.gz` returns `.tar.gz` and `v1.2.zip` returns `.zip`.
+        """
+        last_suffix = path.suffix.lower()
+        compound_exts = cls._compound_exts_by_last_suffix.get(last_suffix)
+        if compound_exts:
+            name = path.name.lower()
+            for ext in compound_exts:
+                if name.endswith(ext):
+                    return ext
+        return last_suffix
+
+    @classmethod
+    def get_compound_exts(cls, ext: str) -> set[str]:
+        """Return the registered compound extensions that start with an extension.
+
+        For example, `.tar` returns `.tar.gz`, `.tar.xz`, etc. (but not `.tar` itself).
+        """
+        return {
+            compound_ext
+            for compound_exts in cls._compound_exts_by_last_suffix.values()
+            for compound_ext in compound_exts
+            if compound_ext.startswith(f"{ext}.")
+        }
+
 
 _FORBIDDEN_NAMES = set(dir(MediaTypes))

src/tagstudio/core/query_lang/file_groups.py

@@ -346,7 +346,7 @@ def register_types() -> None:
     MediaTypes.register("model", ".stl", SEARCH)
 
     # Archives -----------------------------------------------------------------
-    MediaTypes.add_name_aliases("archive", ["Archive", "Compressed"])
+    MediaTypes.add_name_aliases("archive", "Archive")
     MediaTypes.register("archive", ".cba", SEARCH)  # Also under "ebook.comic"
 
     # RAR
@@ -374,7 +374,9 @@ def register_types() -> None:
         ],
     )
     MediaTypes.register("archive.tar", ".tar", SEARCH)
-    MediaTypes.register("archive.tar", [".tar.bz", ".tb2", ".tbz", ".tbz2", ".tz2"], SEARCH)
+    MediaTypes.register(
+        "archive.tar", [".tar.bz", ".tar.bz2", ".tb2", ".tbz", ".tbz2", ".tz2"], SEARCH
+    )
     MediaTypes.register("archive.tar", [".tar.gz", ".taz", ".tgz"], SEARCH)
     MediaTypes.register("archive.tar", [".tar.lzma", ".tlz"], SEARCH)
     MediaTypes.register("archive.tar", [".tar.xz", ".txz"], SEARCH)
@@ -434,6 +436,16 @@ def register_types() -> None:
     MediaTypes.register("binary", ".pyd", SEARCH)
     MediaTypes.register("binary", ".pyo", SEARCH)
 
+    # Compressed ---------------------------------------------------------------
+    # NOTE: These are single-file compression formats that wrap other files (e.g. ".tar.gz").
+    MediaTypes.add_name_aliases("compressed", "Compressed")
+    MediaTypes.register("compressed", ".bz", SEARCH)
+    MediaTypes.register("compressed", ".bz2", SEARCH)
+    MediaTypes.register("compressed", ".gz", SEARCH)
+    MediaTypes.register("compressed", ".lzma", SEARCH)
+    MediaTypes.register("compressed", ".xz", SEARCH)
+    MediaTypes.register("compressed", ".zst", SEARCH)
+
     # Databases ----------------------------------------------------------------
     MediaTypes.add_name_aliases("database", ["Database", "DB"])
     MediaTypes.register("database", ".db", SEARCH)

src/tagstudio/core/utils/decompression.py

@@ -0,0 +1,87 @@
+# SPDX-FileCopyrightText: (c) TagStudio Contributors
+# SPDX-License-Identifier: MIT
+
+
+from collections.abc import Callable, Generator
+from compression import bz2, gzip, lzma, zstd
+from contextlib import contextmanager
+from io import BufferedIOBase
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+import structlog
+
+from tagstudio.core.media_types import MediaTypes
+from tagstudio.core.query_lang.file_groups import SEARCH
+
+logger = structlog.get_logger(__name__)
+
+MAX_DECOMPRESSED_SIZE: int = 128 * 1024 * 1024  # 128 MiB
+_CHUNK_SIZE: int = 1024 * 1024
+
+_OPENERS: dict[str, Callable[[Path], BufferedIOBase]] = {
+    ".bz": lambda path: bz2.open(path, "rb"),
+    ".bz2": lambda path: bz2.open(path, "rb"),
+    ".gz": lambda path: gzip.open(path, "rb"),
+    ".lzma": lambda path: lzma.open(path, "rb"),
+    ".xz": lambda path: lzma.open(path, "rb"),
+    ".zst": lambda path: zstd.open(path, "rb"),
+}
+
+
+def is_compressed(path: Path) -> bool:
+    """Return `True` if the file is a compressed file."""
+    return MediaTypes.contains("compressed", path.suffix.lower(), SEARCH)
+
+
+def get_inner_ext(path: Path) -> str | None:
+    """Return the extension of the file inside a compressed file, if it's a known file type."""
+    if not is_compressed(path) or MediaTypes.get_ext(path) != path.suffix.lower():
+        return None
+    inner_ext = MediaTypes.get_ext(Path(path.stem))
+    return inner_ext if MediaTypes.find(inner_ext, SEARCH) else None
+
+
+def get_display_ext(path: Path) -> str:
+    """Return a file's extension, including the inner file type if compressed (e.g. ".png.gz")."""
+    inner_ext = get_inner_ext(path)
+    return f"{inner_ext}{path.suffix.lower()}" if inner_ext else MediaTypes.get_ext(path)
+
+
+def decompress(path: Path, dest: Path, max_size: int = MAX_DECOMPRESSED_SIZE) -> bool:
+    """Decompress a file to dest, stopping and returning `False` if it exceeds max_size bytes."""
+    size = 0
+    with _OPENERS[path.suffix.lower()](path) as source, open(dest, "wb") as out:
+        while chunk := source.read(_CHUNK_SIZE):
+            size += len(chunk)
+            if size > max_size:
+                return False
+            out.write(chunk)
+    return True
+
+
+@contextmanager
+def decompressed(path: Path, max_size: int = MAX_DECOMPRESSED_SIZE) -> Generator[Path | None]:
+    """Yield a path to a file's uncompressed contents, or `None` if it can't be decompressed."""
+    if not is_compressed(path):
+        yield path
+        return
+
+    if path.suffix.lower() not in _OPENERS:
+        logger.info("[Decompression] Unsupported compression format", path=path)
+        yield None
+        return
+
+    # NOTE: Windows can't delete a file that a reader still has open, so cleanup can't be strict.
+    with TemporaryDirectory(ignore_cleanup_errors=True) as temp_dir:
+        temp_path = Path(temp_dir) / path.stem
+        try:
+            is_decompressed = decompress(path, temp_path, max_size)
+        except (EOFError, OSError, lzma.LZMAError, zstd.ZstdError) as e:
+            logger.error("[Decompression] Could not decompress file", path=path, error=e)
+            is_decompressed = False
+        else:
+            if not is_decompressed:
+                logger.warning("[Decompression] File is too large to decompress", path=path)
+
+        yield temp_path if is_decompressed else None

src/tagstudio/previews/file_renderer.py

@@ -19,6 +19,7 @@
 from tagstudio.core.library.ignore import Ignore
 from tagstudio.core.media_types import MediaTypeGroup, MediaTypes, slugify
 from tagstudio.core.query_lang.file_groups import SEARCH
+from tagstudio.core.utils.decompression import decompressed, get_inner_ext, is_compressed
 from tagstudio.core.utils.types import unwrap
 from tagstudio.previews.base_preview import RENDER, BasePreview
 from tagstudio.previews.effects import apply_overlay_color
@@ -104,15 +105,15 @@ def __init__(self, library: Library, settings: AppSettings) -> None:
         # Key: ("name", UiColor, 512, 512, 1.25)
         self.icons: dict[tuple[str, UiColor, int, int, float], Image.Image] = {}
 
-    def _get_resource_id(self, url: Path) -> str:
+    def _get_resource_id(self, path: Path) -> str:
         """Return the name of the icon resource to use for a file type.
 
         Special terms will return special resources.
 
         Args:
-            url (Path): The file url to assess. "$LOADING" will return the loading graphic.
+            path (Path): The file path to assess. "$LOADING" will return the loading graphic.
         """
-        ext = url.suffix.lower()
+        ext = get_inner_ext(path) or MediaTypes.get_ext(path)
         groups = MediaTypes.find(ext, SEARCH)  # Fallback icons use the SEARCH context
         groups.sort(  # Sort by priority and most specific dot-separated subgroup.
             key=lambda g: (
@@ -731,27 +732,19 @@ def _render(
 
         if filepath and filepath.is_file():
             try:
-                ext = filepath.suffix.lower() if filepath.suffix else filepath.stem.lower()
-                for preview in FileRenderer._preview_renderers:
-                    media_type: MediaTypeGroup | None = getattr(
-                        MediaTypes, preview.media_type_name, None
+                preview = FileRenderer._find_preview(filepath)
+                if preview:
+                    image = preview.render(
+                        filepath=filepath,
+                        is_small=is_small,
+                        theme=theme,
+                        size=(scaled_size, scaled_size),
+                        dpi_scale=dpi_scale,
+                    )
+                elif is_compressed(filepath):
+                    image = FileRenderer._render_decompressed(
+                        filepath, is_small, theme, (scaled_size, scaled_size), dpi_scale
                     )
-                    if media_type is None:
-                        logger.error(
-                            f"[FileRenderer] "
-                            f"Attribute '{preview.media_type_name}' not registered with MediaTypes",
-                        )
-                        break
-
-                    if media_type.contains(ext, RENDER):
-                        image = preview.render(
-                            filepath=filepath,
-                            is_small=is_small,
-                            theme=theme,
-                            size=(scaled_size, scaled_size),
-                            dpi_scale=dpi_scale,
-                        )
-                        break
 
                 if image:
                     image = self._resize_image(image, (scaled_size, scaled_size))
@@ -763,6 +756,44 @@ def _render(
 
         return image
 
+    @staticmethod
+    def _find_preview(filepath: Path) -> type[BasePreview] | None:
+        """Return the highest priority preview renderer for a file's extension, if any."""
+        ext = MediaTypes.get_ext(filepath) if filepath.suffix else filepath.stem.lower()
+        for preview in FileRenderer._preview_renderers:
+            media_type: MediaTypeGroup | None = getattr(MediaTypes, preview.media_type_name, None)
+            if media_type is None:
+                logger.error(
+                    f"[FileRenderer] "
+                    f"Attribute '{preview.media_type_name}' not registered with MediaTypes",
+                )
+                return None
+
+            if media_type.contains(ext, RENDER):
+                return preview
+
+        return None
+
+    @staticmethod
+    def _render_decompressed(
+        filepath: Path, is_small: bool, theme: Theme, size: tuple[int, int], dpi_scale: float
+    ) -> Image.Image | None:
+        """Render a compressed file with the preview renderer for the uncompressed file."""
+        preview = FileRenderer._find_preview(Path(filepath.stem))
+        if preview is None:
+            return None
+
+        with decompressed(filepath) as inner_path:
+            if inner_path is None:
+                return None
+
+            image = preview.render(
+                filepath=inner_path, is_small=is_small, theme=theme, size=size, dpi_scale=dpi_scale
+            )
+            if image:
+                image.load()  # Read the image data before the temporary file is deleted
+            return image
+
     def _resize_image(self, image: Image.Image, size: tuple[int, int]) -> Image.Image:
         orig_x, orig_y = image.size
         new_x, new_y = size