diff --git a/docs/preview-support.md b/docs/preview-support.md index df64c56db..876d005f5 100644 --- a/docs/preview-support.md +++ b/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 diff --git a/src/tagstudio/core/library/alchemy/visitors.py b/src/tagstudio/core/library/alchemy/visitors.py index ca61bd4ab..4290d79b8 100644 --- a/src/tagstudio/core/library/alchemy/visitors.py +++ b/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: diff --git a/src/tagstudio/core/media_types.py b/src/tagstudio/core/media_types.py index 1aa9446c5..f094b8587 100644 --- a/src/tagstudio/core/media_types.py +++ b/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)) diff --git a/src/tagstudio/core/query_lang/file_groups.py b/src/tagstudio/core/query_lang/file_groups.py index c1373031b..546bd7374 100644 --- a/src/tagstudio/core/query_lang/file_groups.py +++ b/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) diff --git a/src/tagstudio/core/utils/decompression.py b/src/tagstudio/core/utils/decompression.py new file mode 100644 index 000000000..dfe864f63 --- /dev/null +++ b/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 diff --git a/src/tagstudio/previews/file_renderer.py b/src/tagstudio/previews/file_renderer.py index 73086bb9a..f89136038 100644 --- a/src/tagstudio/previews/file_renderer.py +++ b/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 diff --git a/src/tagstudio/previews/renderers/archive.py b/src/tagstudio/previews/renderers/archive.py index 800179f5d..1f1b612cf 100644 --- a/src/tagstudio/previews/renderers/archive.py +++ b/src/tagstudio/previews/renderers/archive.py @@ -33,12 +33,16 @@ class ArchivePreview(BasePreview): def register_types(cls) -> None: # NOTE: Filetype equivalents (i.e. ".tar.gz" == ".tgz") are already declared internally. MediaTypes.register("archive", ".7z", RENDER) - MediaTypes.register("archive", ".gz", RENDER) MediaTypes.register("archive", ".rar", RENDER) MediaTypes.register("archive", ".s7z", RENDER) MediaTypes.register("archive", ".tar", RENDER) MediaTypes.register("archive", ".zip", RENDER) + MediaTypes.register("archive", ".tar.bz", RENDER) + MediaTypes.register("archive", ".tar.bz2", RENDER) MediaTypes.register("archive", ".tar.gz", RENDER) + MediaTypes.register("archive", ".tar.lzma", RENDER) + MediaTypes.register("archive", ".tar.xz", RENDER) + MediaTypes.register("archive", ".tar.zst", RENDER) @override @classmethod @@ -95,18 +99,26 @@ def open_archive(filepath: Path) -> Archive: Args: filepath (Path): The path to the archive. - ext (str): The file extension. Returns: Archive: The opened archive. """ - ext = filepath.suffix.lower() + exts = MediaTypes.get_equivalent_exts(MediaTypes.get_ext(filepath)) archiver: type[Archive] = zipfile.ZipFile - if ext in {".7z", ".cb7", ".s7z"}: + if exts & {".7z", ".cb7", ".s7z"}: archiver = SevenZipFile - elif ext in {".cbr", ".rar"}: + elif exts & {".cbr", ".rar"}: archiver = rarfile.RarFile - elif ext in {".cbt", ".tar", ".tgz"}: + elif exts & { + ".cbt", + ".tar", + ".tar.bz", + ".tar.bz2", + ".tar.gz", + ".tar.lzma", + ".tar.xz", + ".tar.zst", + }: archiver = TarFile return archiver(filepath, "r") diff --git a/src/tagstudio/qt/controllers/preview_thumb.py b/src/tagstudio/qt/controllers/preview_thumb.py index 25b5b29a8..de02a90c9 100644 --- a/src/tagstudio/qt/controllers/preview_thumb.py +++ b/src/tagstudio/qt/controllers/preview_thumb.py @@ -17,6 +17,7 @@ from tagstudio.core.library.alchemy.library import Library from tagstudio.core.media_types import MediaTypes from tagstudio.core.query_lang.file_groups import SEARCH +from tagstudio.core.utils.decompression import decompressed from tagstudio.previews.video_tester import is_readable_video from tagstudio.qt.mixed.file_attributes import FileAttributeData from tagstudio.qt.utils.file_opener import open_file @@ -39,37 +40,43 @@ def __init__(self, library: Library, driver: QtDriver): def __get_image_stats(self, filepath: Path) -> FileAttributeData: """Get width and height of an image as dict.""" stats = FileAttributeData() - ext = filepath.suffix.lower() - if filepath.is_dir(): - pass - elif MediaTypes.contains("image.raster.raw", ext, SEARCH): - try: - with rawpy.imread(str(filepath)) as raw: - rgb = raw.postprocess() - image = Image.new("L", (rgb.shape[1], rgb.shape[0]), color="black") - stats.width = image.width - stats.height = image.height - except ( - LibRawIOError, - LibRawFileUnsupportedError, - FileNotFoundError, - ): + with decompressed(filepath) as image_path: + if image_path is None: + return stats + ext = image_path.suffix.lower() + + if image_path.is_dir(): pass - elif MediaTypes.contains("image.raster", ext, SEARCH): - try: - image = Image.open(str(filepath)) - stats.width = image.width - stats.height = image.height - except ( - DecompressionBombError, - FileNotFoundError, - NotImplementedError, - UnidentifiedImageError, - ) as e: - logger.error("[PreviewThumb] Could not get image stats", filepath=filepath, error=e) - elif MediaTypes.contains("image.vector", ext, SEARCH): - pass # TODO + elif MediaTypes.contains("image.raster.raw", ext, SEARCH): + try: + with rawpy.imread(str(image_path)) as raw: + rgb = raw.postprocess() + image = Image.new("L", (rgb.shape[1], rgb.shape[0]), color="black") + stats.width = image.width + stats.height = image.height + except ( + LibRawIOError, + LibRawFileUnsupportedError, + FileNotFoundError, + ): + pass + elif MediaTypes.contains("image.raster", ext, SEARCH): + try: + with Image.open(str(image_path)) as image: + stats.width = image.width + stats.height = image.height + except ( + DecompressionBombError, + FileNotFoundError, + NotImplementedError, + UnidentifiedImageError, + ) as e: + logger.error( + "[PreviewThumb] Could not get image stats", filepath=filepath, error=e + ) + elif MediaTypes.contains("image.vector", ext, SEARCH): + pass # TODO return stats diff --git a/src/tagstudio/qt/mixed/file_attributes.py b/src/tagstudio/qt/mixed/file_attributes.py index cb595bd58..409097392 100644 --- a/src/tagstudio/qt/mixed/file_attributes.py +++ b/src/tagstudio/qt/mixed/file_attributes.py @@ -20,6 +20,7 @@ from tagstudio.core.library.ignore import Ignore from tagstudio.core.media_types import MediaTypes from tagstudio.core.query_lang.file_groups import SEARCH +from tagstudio.core.utils.decompression import get_display_ext from tagstudio.core.utils.str_formatting import format_duration from tagstudio.core.utils.types import unwrap from tagstudio.i18n.translations import Translations @@ -139,7 +140,7 @@ def update_stats(self, filepath: Path | None = None, stats: FileAttributeData | self.dimensions_label.setText("") self.dimensions_label.setHidden(True) else: - ext = filepath.suffix.lower() + ext = get_display_ext(filepath) display_path = filepath if self.driver.settings.show_filepath == ShowFilepathOption.SHOW_FULL_PATHS: display_path = filepath diff --git a/src/tagstudio/qt/mixed/item_thumb.py b/src/tagstudio/qt/mixed/item_thumb.py index f789d7798..cec05ed87 100644 --- a/src/tagstudio/qt/mixed/item_thumb.py +++ b/src/tagstudio/qt/mixed/item_thumb.py @@ -17,6 +17,7 @@ from tagstudio.core.library.alchemy.library import Library from tagstudio.core.media_types import MediaTypes from tagstudio.core.query_lang.file_groups import SEARCH +from tagstudio.core.utils.decompression import get_display_ext from tagstudio.core.utils.types import unwrap from tagstudio.i18n.platform_strings import open_file_str, trash_term from tagstudio.i18n.translations import Translations @@ -365,7 +366,7 @@ def set_extension(self, filename: Path) -> None: show_ext_badge = False show_count_badge = False - ext = filename.suffix.lower() + ext = get_display_ext(filename) if ext and ext.startswith(".") is False: ext = "." + ext if ( diff --git a/tests/core/library/test_library.py b/tests/core/library/test_library.py index e260225a8..a6cb24ab8 100644 --- a/tests/core/library/test_library.py +++ b/tests/core/library/test_library.py @@ -480,3 +480,54 @@ def test_filetype_return_one_filetype( def test_mediatype_search(library: Library, mediatype: str, num_of_mediatype: int): results = library.search_library(BrowsingState.from_mediatype(mediatype), page_size=500) assert len(results.ids) == num_of_mediatype + + +@pytest.mark.parametrize( + ["query", "paths"], + [ + ("filetype:tar.gz", {"backup.tar.gz", "photos.tgz"}), + ("filetype:TAR.GZ", {"backup.tar.gz", "photos.tgz"}), + ("filetype:tgz", {"backup.tar.gz", "photos.tgz"}), + ("filetype:tar.xz", {"my.cool.archive.tar.xz", "CAPITAL.TAR.XZ"}), + ( + "filetype:tar", + { + "plain.tar", + "backup.tar.gz", + "photos.tgz", + "my.cool.archive.tar.xz", + "CAPITAL.TAR.XZ", + }, + ), + ( + "mediatype:archive", + { + "plain.tar", + "backup.tar.gz", + "photos.tgz", + "my.cool.archive.tar.xz", + "CAPITAL.TAR.XZ", + }, + ), + ( + "mediatype:compressed", + {"backup.tar.gz", "my.cool.archive.tar.xz", "CAPITAL.TAR.XZ", "log.gz"}, + ), + ], +) +def test_compound_ext_search(query: str, paths: set[str]): + lib = Library() + assert lib.open_library(Path(""), in_memory=True).success + names = [ + "plain.tar", + "backup.tar.gz", + "photos.tgz", + "my.cool.archive.tar.xz", + "CAPITAL.TAR.XZ", + "log.gz", + "notes.txt", + ] + assert lib.add_entries([Entry(path=Path(name), fields=[]) for name in names]) + + results = lib.search_library(BrowsingState.from_search_query(query), page_size=500) + assert {str(unwrap(lib.get_entry(entry_id)).path) for entry_id in results.ids} == paths diff --git a/tests/core/test_decompression.py b/tests/core/test_decompression.py new file mode 100644 index 000000000..2ce823dd1 --- /dev/null +++ b/tests/core/test_decompression.py @@ -0,0 +1,103 @@ +# SPDX-FileCopyrightText: (c) TagStudio Contributors +# SPDX-License-Identifier: MIT + +# pyright: reportPrivateUsage=false + + +from compression import gzip +from pathlib import Path + +import pytest + +from tagstudio.core.media_types import MediaTypes +from tagstudio.core.query_lang.file_groups import SEARCH +from tagstudio.core.utils.decompression import ( + decompress, + decompressed, + get_display_ext, + get_inner_ext, + is_compressed, +) + + +@pytest.mark.parametrize( + ["filename", "expected"], + [("photo.png.gz", True), ("Photo.PNG.XZ", True), ("photo.png", False), ("bundle.zip", False)], +) +def test_is_compressed(filename: str, expected: bool): + assert is_compressed(Path(filename)) is expected + + +def test_unsupported_compression_format(tmp_path: Path): + MediaTypes.register("compressed", ".zzzcomp", SEARCH) + source = tmp_path / "photo.png.zzzcomp" + source.touch() + + assert is_compressed(source) + assert get_display_ext(source) == ".png.zzzcomp" + with decompressed(source) as path: + assert path is None + + +@pytest.mark.parametrize( + ["filename", "inner_ext", "display_ext"], + [ + ("photo.png.gz", ".png", ".png.gz"), + ("my.cool.Photo.PNG.XZ", ".png", ".png.xz"), + ("bundle.tar.gz", None, ".tar.gz"), + ("notes.final.gz", None, ".gz"), + ("log.gz", None, ".gz"), + ("photo.png", None, ".png"), + ], +) +def test_inner_and_display_ext(filename: str, inner_ext: str | None, display_ext: str): + assert get_inner_ext(Path(filename)) == inner_ext + assert get_display_ext(Path(filename)) == display_ext + + +def test_decompress(tmp_path: Path): + source = tmp_path / "notes.txt.gz" + source.write_bytes(gzip.compress(b"x" * 1000)) # Decompresses to exactly 1000 bytes + + # A size equal to the limit is allowed (1000) + assert decompress(source, tmp_path / "notes.txt", max_size=1000) + # The full contents were written + assert (tmp_path / "notes.txt").read_bytes() == b"x" * 1000 + # A size over the limit is not allowed (999) + assert not decompress(source, tmp_path / "capped.txt", max_size=999) + + +def test_decompressed_uncompressed_file(tmp_path: Path): + source = tmp_path / "notes.txt" + source.write_bytes(b"plain") + + with decompressed(source) as path: + assert path == source + + +def test_decompressed_compressed_file(tmp_path: Path): + source = tmp_path / "my.notes.txt.gz" + source.write_bytes(gzip.compress(b"x" * 1000)) + + with decompressed(source) as path: + assert path is not None + assert path.name == "my.notes.txt" + assert path.read_bytes() == b"x" * 1000 + + assert not path.exists() + + +def test_decompressed_too_large(tmp_path: Path): + source = tmp_path / "notes.txt.gz" + source.write_bytes(gzip.compress(b"x" * 1000)) + + with decompressed(source, max_size=999) as path: + assert path is None + + +def test_decompressed_corrupt_file(tmp_path: Path): + source = tmp_path / "photo.png.gz" + source.write_bytes(b"not gzip data") + + with decompressed(source) as path: + assert path is None diff --git a/tests/core/test_media_types.py b/tests/core/test_media_types.py index 307cf1f7f..c62d0bbed 100644 --- a/tests/core/test_media_types.py +++ b/tests/core/test_media_types.py @@ -2,6 +2,8 @@ # SPDX-License-Identifier: MIT +from pathlib import Path + from pytestqt.exceptions import pytest from tagstudio.core.media_types import MediaTypes @@ -55,6 +57,22 @@ def test_get_equivalent_exts_defaults_to_itself(): assert MediaTypes.get_equivalent_exts(".zzzunregistered") == {".zzzunregistered"} +@pytest.mark.parametrize( + ["filename", "ext"], + [ + ("pack.zzzinner.zzzouter", ".zzzinner.zzzouter"), + ("my.dotted.name.ZZZINNER.ZZZOUTER", ".zzzinner.zzzouter"), + ("pack.zzzouter", ".zzzouter"), + ("pack.other.zzzouter", ".zzzouter"), + ("README", ""), + ], +) +def test_get_ext(filename: str, ext: str): + MediaTypes.register("zzztest.compound", ".zzzinner.zzzouter", "SEARCH") + + assert MediaTypes.get_ext(Path(filename)) == ext + + def test_find(): MediaTypes.register("zzztest.find_a", ".zzzfind", "SEARCH") MediaTypes.register("zzztest.find_b", ".zzzfind", "RENDER") diff --git a/tests/previews/test_file_renderer.py b/tests/previews/test_file_renderer.py new file mode 100644 index 000000000..7d03adb51 --- /dev/null +++ b/tests/previews/test_file_renderer.py @@ -0,0 +1,90 @@ +# SPDX-FileCopyrightText: (c) TagStudio Contributors +# SPDX-License-Identifier: MIT + +# pyright: reportPrivateUsage=false + + +import tarfile +from collections.abc import Callable +from compression import bz2, gzip, lzma, zstd +from io import BytesIO +from pathlib import Path + +import pytest +from PIL import Image + +from tagstudio.core.enums import Theme +from tagstudio.previews.file_renderer import FileRenderer +from tagstudio.previews.renderers.archive import ArchivePreview + + +def _example_png_bytes() -> bytes: + data = BytesIO() + Image.new("RGBA", (40, 20), (255, 0, 0, 255)).save(data, "PNG") + return data.getvalue() + + +@pytest.mark.parametrize("filename", ["photo.png.gz", "my.cool.photo.PNG.XZ"]) +def test_render_decompressed(tmp_path: Path, filename: str): + source = tmp_path / filename + if filename.lower().endswith(".gz"): + source.write_bytes(gzip.compress(_example_png_bytes())) + else: + source.write_bytes(lzma.compress(_example_png_bytes())) + + image = FileRenderer._render_decompressed( + source, is_small=True, theme=Theme.DARK, size=(128, 128), dpi_scale=1.0 + ) + + assert image is not None + assert image.size == (40, 20) + + +def _example_tar_bytes() -> bytes: + data = BytesIO() + with tarfile.open(fileobj=data, mode="w") as tar: + png = _example_png_bytes() + info = tarfile.TarInfo("images/photo.png") + info.size = len(png) + tar.addfile(info, BytesIO(png)) + return data.getvalue() + + +@pytest.mark.parametrize( + ["filename", "compress"], + [ + ("bundle.tar.bz2", bz2.compress), + ("bundle.tbz2", bz2.compress), + ("bundle.tar.lzma", lambda data: lzma.compress(data, format=lzma.FORMAT_ALONE)), + ("bundle.tlz", lambda data: lzma.compress(data, format=lzma.FORMAT_ALONE)), + ("my.cool.bundle.tar.xz", lzma.compress), + ("BUNDLE.TXZ", lzma.compress), + ("bundle.tar.zst", zstd.compress), + ("bundle.tzst", zstd.compress), + ], +) +def test_render_compressed_tarball( + tmp_path: Path, filename: str, compress: Callable[[bytes], bytes] +): + source = tmp_path / filename + source.write_bytes(compress(_example_tar_bytes())) + + assert FileRenderer._find_preview(source) is ArchivePreview + image = ArchivePreview.render( + source, is_small=True, theme=Theme.DARK, size=(128, 128), dpi_scale=1.0 + ) + + assert image is not None + assert image.size == (40, 20) + + +def test_render_decompressed_without_inner_extension(tmp_path: Path): + source = tmp_path / "log.gz" + source.write_bytes(gzip.compress(b"plain text")) + + assert ( + FileRenderer._render_decompressed( + source, is_small=True, theme=Theme.DARK, size=(128, 128), dpi_scale=1.0 + ) + is None + )