From ba23a185be1872e4e621e0107dbebfae0b8d4b42 Mon Sep 17 00:00:00 2001 From: Abdussamet Kocak Date: Tue, 20 Oct 2026 07:29:03 +0300 Subject: [PATCH] fix(download_imageset): Skip processed images --- download_imageset.py | 207 ++++++++++++++++++++--------------- download_imageset_test.py | 219 +++++++++++++++++++++++++++++++++++++- 2 files changed, 336 insertions(+), 90 deletions(-) diff --git a/download_imageset.py b/download_imageset.py index fe35ab5..8516c96 100755 --- a/download_imageset.py +++ b/download_imageset.py @@ -7,6 +7,7 @@ import dataclasses import logging import os import re +import shutil import struct import subprocess import urllib.parse @@ -17,7 +18,7 @@ from typing import Callable import click import httpx -from PIL import Image +from PIL import Image, ImageOps run_command = partial( subprocess.run, @@ -201,12 +202,21 @@ def _has_images(d: Path) -> bool: def _needs_resize(d: Path) -> bool: - images = [f for f in d.rglob("*") if f.is_file() and f.suffix.lower() in IMAGE_EXTS] + images = [f for f in d.rglob("*") if f.is_file() and f.suffix.lower() in IMAGE_EXTS and not f.is_relative_to(d / "resized")] if not images: return False mean_mb = sum(f.stat().st_size for f in images) / len(images) files_are_large = mean_mb >= 2 * 1_048_567 - images_are_large = any(max(read_dims(img)) > 5000 for img in images) + images_are_large = False + for image in images: + try: + if max(read_dims(image)) > 5000: + images_are_large = True + break + except Exception as error: + logging.warning(f"could not read dimensions of {image}; trying resizer: {error}") + images_are_large = True + break return files_are_large or images_are_large @@ -216,46 +226,111 @@ def _needs_resize(d: Path) -> bool: # --------------------------------------------------------------------------- -def _extract_and_process(output_dir: Path) -> None: - archives = [ - f for pattern in ("*.zip", "*.rar", "*.7z") for f in output_dir.glob(pattern) - ] - for archive in archives: - print(f"extracting {archive.name} …") - run_command(["unar", "-o", str(output_dir), str(archive)]) - move_to_trash([archive]) +ImageProcessor = Callable[..., None] - subdirs = [d for d in output_dir.iterdir() if d.is_dir() and d.name != "resized"] - targets = subdirs if subdirs else [output_dir] - to_process = [d for d in targets if _needs_resize(d)] - if not to_process: - print("no dirs need resizing") +def detect_resizer() -> ImageProcessor: + vips_binary = shutil.which(cmd="vips") + if vips_binary: + return partial(resize_with_vips, vips_binary=vips_binary) + return resize_with_pillow + + +def resize_with_vips(image_path: Path, output_path: Path, vips_binary: str) -> None: + try: + subprocess.run( + args=[ + vips_binary, + "thumbnail", + str(image_path), + f"{output_path}[Q=75]", + "5000", + "--height", + "5000", + "--size", + "down", + ], + check=True, + capture_output=True, + ) + except (OSError, subprocess.CalledProcessError) as error: + logging.warning(f"vips failed for {image_path}; using Pillow: {error}") + resize_with_pillow(image_path=image_path, output_path=output_path) + + +def resize_with_pillow(image_path: Path, output_path: Path) -> None: + with Image.open(fp=image_path) as image: + with ImageOps.exif_transpose(image=image) as oriented: + oriented.thumbnail(size=(5000, 5000), resample=Image.Resampling.LANCZOS) + with oriented.convert(mode="RGB") as resized: + resized.save(fp=output_path, format="JPEG", quality=75) + + +def process_imageset(path: Path, processor: ImageProcessor) -> None: + if not path.exists(): + logging.warning(f"skipping {path}: does not exist") + return + if path.is_file() and path.suffix.lower() in {".zip", ".rar", ".7z"}: + output_dir = path.with_suffix("") + archives = [path] + elif path.is_dir(): + output_dir = path + archives = [f for f in path.iterdir() if f.is_file() and f.suffix.lower() in {".zip", ".rar", ".7z"}] else: + logging.warning(f"skipping {path}: expected a directory or archive") + return + if output_dir.exists() and not output_dir.is_dir(): + logging.warning(f"skipping {path}: {output_dir} is not a directory") + return + + for archive in archives: + logging.info(f"extracting {archive} → {output_dir}") + if archive.suffix.lower() == ".zip": + try: + with zipfile.ZipFile(file=archive) as zipped: + zipped.extractall(path=output_dir) + except zipfile.BadZipFile: + logging.warning(f"skipping {archive}: invalid ZIP file") + return + else: + run_command(args=["unar", "--output-directory", str(output_dir), str(archive)]) + + output_dir.mkdir(parents=True, exist_ok=True) + failed_images: set[Path] = set() + if _needs_resize(output_dir): resized_dir = output_dir / "resized" resized_dir.mkdir(exist_ok=True) - originals = [ - f - for d in to_process - for f in d.rglob("*") - if f.is_file() and f.suffix.lower() in IMAGE_EXTS - ] - print(f"resizing {len(originals)} image(s) → {resized_dir}") + originals = [f for f in output_dir.rglob("*") if is_image(f) and not f.is_relative_to(resized_dir)] + logging.info(f"resizing {len(originals)} image(s) → {resized_dir}") + resized_originals: list[Path] = [] for img in originals: - run_command( - [ - "imgz", - "resize", - "-o", - str(resized_dir / f"{img.stem}.resized.jpg"), - str(img), - ] - ) - if originals: - print(f"trashing {len(originals)} original(s)") - move_to_trash(originals) - rename_children(output_dir, flatten=True) - remove_empty_dirs(output_dir) + output_path = resized_dir / f"{img.stem}.resized.jpg" + if output_path.is_file() and output_path.stat().st_size > 0: + logging.info(f"skipping {img}: {output_path} already resized") + resized_originals.append(img) + continue + try: + processor(image_path=img, output_path=output_path) + except Exception as error: + logging.warning(f"skipping {img}: resizing failed: {error}") + failed_images.add(img) + output_path.unlink(missing_ok=True) + continue + resized_originals.append(img) + if resized_originals: + logging.info(f"trashing {len(resized_originals)} original(s)") + move_to_trash(file_paths=resized_originals) + else: + logging.info(f"skipping resizing {output_dir}: already within size limits") + + rename_children( + source_dir=output_dir, + flatten=True, + filter_fn=lambda image: is_image(image) and image not in failed_images, + ) + remove_empty_dirs(root_dir=output_dir) + if archives and not failed_images: + move_to_trash(file_paths=archives) def remove_empty_dirs(root_dir: Path) -> None: @@ -356,7 +431,7 @@ def gallery_cmd(urls: tuple[str, ...], name: str | None, cwd: Path | None) -> No args += ["-D", str(output_dir)] run_command(args + urls) - _extract_and_process(output_dir) + process_imageset(path=output_dir, processor=detect_resizer()) @cli.command("filehost") @@ -389,10 +464,10 @@ def filehost_cmd(urls: tuple[str, ...], name: str | None, cwd: Path, token: str) print(f"downloading {meta.suggested_filename} → {output_dir}") run_command(["aria2c", f"--dir={output_dir}", dl_url]) - _extract_and_process(output_dir) + process_imageset(path=output_dir, processor=detect_resizer()) def is_filehost(url: str) -> bool: - kws = ["rg.to", "rapidgator.net", "trbt.cc"] + kws = ["rg.to", "rapidgator.net", "trbt.cc", "katfile"] return any(kw in url for kw in kws) @cli.command("download") @@ -415,54 +490,10 @@ def download_cmd( @cli.command("process") @click.argument("paths", nargs=-1, required=True, type=click.Path(path_type=Path)) def process_cmd(paths: tuple[Path, ...]) -> None: - """Process directories and ZIP files, extracting ZIPs into sibling directories.""" - for p in paths: - archive_path: Path | None = None - if not p.exists(): - logging.warning(f"skipping {p}: does not exist") - continue - if p.is_file() and p.suffix.lower() == ".zip": - output_dir = p.with_suffix("") - if output_dir.exists() and not output_dir.is_dir(): - logging.warning(f"skipping {p}: {output_dir} is not a directory") - continue - try: - with zipfile.ZipFile(file=p) as archive: - logging.info(f"extracting {p} → {output_dir}") - archive.extractall(path=output_dir) - except zipfile.BadZipFile: - logging.warning(f"skipping {p}: invalid ZIP file") - continue - archive_path = p - p = output_dir - if not p.is_dir(): - logging.warning(f"skipping {p}: expected a directory or ZIP file") - continue - if _needs_resize(p): - resized_dir = p / "resized" - resized_dir.mkdir(exist_ok=True) - originals = [f for f in p.rglob("*") if is_image(f)] - logging.info(f"resizing {len(originals)} image(s) in {p} → {resized_dir}") - for img in originals: - run_command( - [ - "imgz", - "resize", - "--output", - str(resized_dir / f"{img.stem}.resized.jpg"), - str(img), - ] - ) - if originals: - logging.info(f"trashing {len(originals)} original(s)") - move_to_trash(originals) - else: - logging.info(f"skipping {p}: already within size limits") - - rename_children(p, flatten=True) - remove_empty_dirs(p) - if archive_path is not None: - move_to_trash(file_paths=[archive_path]) + """Process directories and archives, extracting archives into sibling directories.""" + processor = detect_resizer() + for path in paths: + process_imageset(path=path, processor=processor) if __name__ == "__main__": diff --git a/download_imageset_test.py b/download_imageset_test.py index c22da76..56121c4 100644 --- a/download_imageset_test.py +++ b/download_imageset_test.py @@ -3,9 +3,12 @@ # requires-python = ">=3.11" # dependencies = ["pytest", "httpx", "pillow", "click"] # /// +import shutil +import subprocess import sys import zipfile from pathlib import Path +from functools import partial from unittest.mock import Mock import pytest @@ -43,6 +46,47 @@ def test_process_zip_and_directory( trash.assert_called_once_with(file_paths=[archive_path]) +@pytest.mark.parametrize("command", ["process", "gallery"]) +@pytest.mark.parametrize("processing_fails", [False, True]) +def test_process_archives_in_directory( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + command: str, + processing_fails: bool, +) -> None: + directory = tmp_path / "album" + directory.mkdir() + archive_path = directory / "download.ZIP" + image_path = tmp_path / "photo.png" + Image.new(mode="RGB", size=(10, 10)).save(fp=image_path) + with zipfile.ZipFile(file=archive_path, mode="w") as archive: + if processing_fails: + archive.writestr(zinfo_or_arcname="nested/photo.png", data=b"invalid image") + else: + archive.write(filename=image_path, arcname="nested/photo.png") + trash = Mock() + monkeypatch.setattr(download_imageset, "move_to_trash", trash) + monkeypatch.setattr(download_imageset, "run_command", Mock()) + args = ( + ["process", str(directory)] + if command == "process" + else ["gallery", "https://example.com/album", "--cwd", str(tmp_path)] + ) + + result = CliRunner().invoke(cli=download_imageset.cli, args=args) + + if processing_fails: + assert result.exit_code == 0, result.output + assert archive_path.is_file() + assert (directory / "nested" / "photo.png").read_bytes() == b"invalid image" + trash.assert_not_called() + else: + assert result.exit_code == 0, result.output + assert (directory / "album__0000.png").is_file() + assert not (directory / "nested").exists() + trash.assert_called_once_with(file_paths=[archive_path]) + + @pytest.mark.parametrize("filename", ["missing", "unsupported.txt", "broken.zip"]) def test_process_skips_invalid_inputs( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, filename: str @@ -65,7 +109,7 @@ def test_process_skips_invalid_inputs( def test_process_keeps_zip_when_processing_fails( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture ) -> None: archive_path = tmp_path / "gallery.zip" with zipfile.ZipFile(file=archive_path, mode="w") as archive: @@ -77,9 +121,180 @@ def test_process_keeps_zip_when_processing_fails( cli=download_imageset.cli, args=["process", str(archive_path)] ) - assert result.exit_code != 0 + assert result.exit_code == 0, result.output assert archive_path.is_file() trash.assert_not_called() + assert "resizing failed" in caplog.text + + +def test_process_continues_after_resize_failure( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + directory = tmp_path / "album" + directory.mkdir() + failed_image = directory / "broken.png" + failed_image.write_bytes(data=b"broken image") + good_image = directory / "good.png" + Image.new(mode="RGB", size=(6000, 60)).save(fp=good_image) + + def processor(image_path: Path, output_path: Path) -> None: + if image_path == failed_image: + output_path.write_bytes(data=b"partial JPEG") + raise OSError("cannot decode image") + download_imageset.resize_with_pillow(image_path=image_path, output_path=output_path) + + trash = Mock() + monkeypatch.setattr(download_imageset, "detect_resizer", lambda: processor) + monkeypatch.setattr(download_imageset, "move_to_trash", trash) + + result = CliRunner().invoke(cli=download_imageset.cli, args=["process", str(directory)]) + + assert result.exit_code == 0, result.output + assert failed_image.read_bytes() == b"broken image" + assert not (directory / "resized" / "broken.resized.jpg").exists() + with Image.open(fp=directory / "album__0001.jpg") as resized: + assert resized.size == (5000, 50) + trash.assert_called_once_with(file_paths=[good_image]) + assert f"skipping {failed_image}: resizing failed: cannot decode image" in caplog.text + + +@pytest.mark.parametrize("use_vips", [False, True]) +@pytest.mark.parametrize( + ("size", "mode", "orientation", "expected_size"), + [ + ((6000, 60), "RGB", 1, (5000, 50)), + ((60, 6000), "RGB", 1, (50, 5000)), + ((40, 20), "RGBA", 1, (40, 20)), + ((40, 20), "RGB", 6, (20, 40)), + ], +) +def test_resize_image( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + use_vips: bool, + size: tuple[int, int], + mode: str, + orientation: int, + expected_size: tuple[int, int], +) -> None: + vips_binary = shutil.which(cmd="vips") if use_vips else None + if use_vips and vips_binary is None: + pytest.skip("vips is not installed") + processor = ( + partial(download_imageset.resize_with_vips, vips_binary=vips_binary) + if vips_binary + else download_imageset.resize_with_pillow + ) + image_path = tmp_path / "source.png" + output_path = tmp_path / "resized.jpg" + with Image.new(mode=mode, size=size) as image: + exif = Image.Exif() + exif[274] = orientation + image.save(fp=image_path, exif=exif) + + processor(image_path=image_path, output_path=output_path) + + with Image.open(fp=output_path) as resized: + resized.load() + assert resized.format == "JPEG" + assert resized.mode == "RGB" + assert resized.size == expected_size + assert image_path.is_file() + + +@pytest.mark.parametrize( + "error", [OSError("vips unavailable"), subprocess.CalledProcessError(1, "vips")] +) +def test_resize_image_falls_back_when_vips_fails( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, error: Exception +) -> None: + image_path = tmp_path / "source.png" + output_path = tmp_path / "resized.jpg" + Image.new(mode="RGB", size=(6000, 60)).save(fp=image_path) + monkeypatch.setattr(download_imageset.subprocess, "run", Mock(side_effect=error)) + + download_imageset.resize_with_vips( + image_path=image_path, output_path=output_path, vips_binary="vips" + ) + + with Image.open(fp=output_path) as resized: + resized.load() + assert resized.size == (5000, 50) + + +@pytest.mark.parametrize("vips_binary", [None, "/usr/bin/vips"]) +def test_process_selects_processor_once( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, vips_binary: str | None +) -> None: + directories = [tmp_path / "first", tmp_path / "second"] + for directory in directories: + directory.mkdir() + Image.new(mode="RGB", size=(10, 10)).save(fp=directory / "photo.png") + lookup = Mock(return_value=vips_binary) + vips_processor = Mock() + pillow_processor = Mock() + monkeypatch.setattr(download_imageset.shutil, "which", lookup) + monkeypatch.setattr(download_imageset, "resize_with_vips", vips_processor) + monkeypatch.setattr(download_imageset, "resize_with_pillow", pillow_processor) + monkeypatch.setattr(download_imageset, "_needs_resize", Mock(return_value=True)) + monkeypatch.setattr(download_imageset, "move_to_trash", Mock()) + + result = CliRunner().invoke( + cli=download_imageset.cli, args=["process", *map(str, directories)] + ) + + assert result.exit_code == 0, result.output + lookup.assert_called_once_with(cmd="vips") + selected_processor = vips_processor if vips_binary else pillow_processor + unused_processor = pillow_processor if vips_binary else vips_processor + assert selected_processor.call_count == 2 + unused_processor.assert_not_called() + for directory, call in zip(directories, selected_processor.call_args_list): + expected = { + "image_path": directory / "photo.png", + "output_path": directory / "resized" / "photo.resized.jpg", + } + if vips_binary: + expected["vips_binary"] = vips_binary + assert call.kwargs == expected + + +@pytest.mark.parametrize("command", ["process", "gallery"]) +@pytest.mark.parametrize("existing_output", [None, b"", b"already resized"]) +def test_resume_resizing( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + command: str, + existing_output: bytes | None, +) -> None: + directory = tmp_path / "album" + resized_dir = directory / "resized" + resized_dir.mkdir(parents=True) + image_path = directory / "photo.png" + Image.new(mode="RGB", size=(6000, 60)).save(fp=image_path) + output_path = resized_dir / "photo.resized.jpg" + if existing_output is not None: + output_path.write_bytes(data=existing_output) + processor = Mock(wraps=download_imageset.resize_with_pillow) + trash = Mock() + monkeypatch.setattr(download_imageset, "detect_resizer", lambda: processor) + monkeypatch.setattr(download_imageset, "move_to_trash", trash) + if command == "process": + args = ["process", str(directory)] + else: + args = ["gallery", "https://example.com/album", "--cwd", str(tmp_path)] + + result = CliRunner().invoke(cli=download_imageset.cli, args=args) + + assert result.exit_code == 0, result.output + if existing_output: + processor.assert_not_called() + assert (directory / "album__0001.jpg").read_bytes() == existing_output + else: + processor.assert_called_once_with(image_path=image_path, output_path=output_path) + with Image.open(fp=directory / "album__0001.jpg") as resized: + assert resized.size == (5000, 50) + trash.assert_called_once_with(file_paths=[image_path]) if __name__ == "__main__":