fix(download_imageset): Skip processed images

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__":

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__":