abdus/playground
1
0 Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(download_imageset): Skip processed images

Browse Source
This commit is contained in:
abdus committed 2026-10-20 07:29:03 +03:00
1 parent a443eeb042
commit ba23a185be
2 files changed
+335 -89
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

download_imageset.py
+118 -87
View File
@@ -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),
]
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,
)
if originals:
print(f"trashing {len(originals)} original(s)")
move_to_trash(originals)
rename_children(output_dir, flatten=True)
remove_empty_dirs(output_dir)
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
+217 -2
View File
@@ -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__":