diff --git a/download_imageset.py b/download_imageset.py index 6c804be..fe35ab5 100755 --- a/download_imageset.py +++ b/download_imageset.py @@ -9,8 +9,8 @@ import os import re import struct import subprocess -import sys import urllib.parse +import zipfile from functools import partial from pathlib import Path from typing import Callable @@ -413,37 +413,58 @@ def download_cmd( @cli.command("process") -@click.argument("dirs", nargs=-1, required=True) -def process_cmd(dirs: tuple[str, ...]) -> None: - for d in dirs: - p = Path(d) +@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(): - print(f"skipping {p}: does not exist", file=sys.stderr) + 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)] - print(f"resizing {len(originals)} image(s) in {p} → {resized_dir}") + logging.info(f"resizing {len(originals)} image(s) in {p} → {resized_dir}") for img in originals: run_command( [ "imgz", "resize", - "-o", + "--output", str(resized_dir / f"{img.stem}.resized.jpg"), str(img), ] ) if originals: - print(f"trashing {len(originals)} original(s)") + logging.info(f"trashing {len(originals)} original(s)") move_to_trash(originals) else: - print(f"skipping {p}: already within size limits") + 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]) if __name__ == "__main__": + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s") cli() diff --git a/download_imageset_test.py b/download_imageset_test.py new file mode 100644 index 0000000..c22da76 --- /dev/null +++ b/download_imageset_test.py @@ -0,0 +1,86 @@ +#!/usr/bin/env -S uv run --script +# /// script +# requires-python = ">=3.11" +# dependencies = ["pytest", "httpx", "pillow", "click"] +# /// +import sys +import zipfile +from pathlib import Path +from unittest.mock import Mock + +import pytest +from click.testing import CliRunner +from PIL import Image + +import download_imageset + + +@pytest.mark.parametrize("extension", [".zip", ".ZIP"]) +def test_process_zip_and_directory( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, extension: str +) -> None: + directory = tmp_path / "existing" + nested = directory / "nested" + nested.mkdir(parents=True) + image_path = nested / "photo.png" + Image.new(mode="RGB", size=(10, 10)).save(fp=image_path) + archive_path = tmp_path / f"gallery{extension}" + with zipfile.ZipFile(file=archive_path, mode="w") as archive: + archive.write(filename=image_path, arcname="album/photo.png") + trash = Mock() + monkeypatch.setattr(download_imageset, "move_to_trash", trash) + + result = CliRunner().invoke( + cli=download_imageset.cli, + args=["process", str(archive_path), str(directory)], + ) + + assert result.exit_code == 0, result.output + assert (tmp_path / "gallery" / "gallery__0000.png").is_file() + assert (directory / "existing__0000.png").is_file() + assert not (tmp_path / "gallery" / "album").exists() + assert not 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 +) -> None: + source = tmp_path / filename + if filename != "missing": + source.write_text(data="not an archive") + trash = Mock() + monkeypatch.setattr(download_imageset, "move_to_trash", trash) + + result = CliRunner().invoke( + cli=download_imageset.cli, args=["process", str(source)] + ) + + assert result.exit_code == 0, result.output + trash.assert_not_called() + if filename != "missing": + assert source.read_text() == "not an archive" + assert not (tmp_path / "broken").exists() + + +def test_process_keeps_zip_when_processing_fails( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + archive_path = tmp_path / "gallery.zip" + with zipfile.ZipFile(file=archive_path, mode="w") as archive: + archive.writestr(zinfo_or_arcname="photo.png", data=b"invalid image") + trash = Mock() + monkeypatch.setattr(download_imageset, "move_to_trash", trash) + + result = CliRunner().invoke( + cli=download_imageset.cli, args=["process", str(archive_path)] + ) + + assert result.exit_code != 0 + assert archive_path.is_file() + trash.assert_not_called() + + +if __name__ == "__main__": + sys.exit(pytest.main(args=[str(Path(__file__).resolve()), *sys.argv[1:]]))