feat(download_imageset): process ZIP archives

This commit is contained in:
abdus committed 2026-10-08 11:27:24 +00:00
1 parent 1939ef0c0d
commit d572ca199f
2 files changed
+117 -10

No files matched your search

+31 -10
View File
@@ -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()
+86
View File
@@ -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:]]))