fix(download_imageset): Skip processed images

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

No files matched your search

+118 -87
View File
@@ -7,6 +7,7 @@ import dataclasses
import logging import logging
import os import os
import re import re
import shutil
import struct import struct
import subprocess import subprocess
import urllib.parse import urllib.parse
@@ -17,7 +18,7 @@ from typing import Callable
import click import click
import httpx import httpx
from PIL import Image from PIL import Image, ImageOps
run_command = partial( run_command = partial(
subprocess.run, subprocess.run,
@@ -201,12 +202,21 @@ def _has_images(d: Path) -> bool:
def _needs_resize(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: if not images:
return False return False
mean_mb = sum(f.stat().st_size for f in images) / len(images) mean_mb = sum(f.stat().st_size for f in images) / len(images)
files_are_large = mean_mb >= 2 * 1_048_567 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 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: ImageProcessor = Callable[..., 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])
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)] def detect_resizer() -> ImageProcessor:
if not to_process: vips_binary = shutil.which(cmd="vips")
print("no dirs need resizing") 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: 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 = output_dir / "resized"
resized_dir.mkdir(exist_ok=True) resized_dir.mkdir(exist_ok=True)
originals = [ originals = [f for f in output_dir.rglob("*") if is_image(f) and not f.is_relative_to(resized_dir)]
f logging.info(f"resizing {len(originals)} image(s) → {resized_dir}")
for d in to_process resized_originals: list[Path] = []
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}")
for img in originals: for img in originals:
run_command( output_path = resized_dir / f"{img.stem}.resized.jpg"
[ if output_path.is_file() and output_path.stat().st_size > 0:
"imgz", logging.info(f"skipping {img}: {output_path} already resized")
"resize", resized_originals.append(img)
"-o", continue
str(resized_dir / f"{img.stem}.resized.jpg"), try:
str(img), 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: remove_empty_dirs(root_dir=output_dir)
print(f"trashing {len(originals)} original(s)") if archives and not failed_images:
move_to_trash(originals) move_to_trash(file_paths=archives)
rename_children(output_dir, flatten=True)
remove_empty_dirs(output_dir)
def remove_empty_dirs(root_dir: Path) -> None: 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)] args += ["-D", str(output_dir)]
run_command(args + urls) run_command(args + urls)
_extract_and_process(output_dir) process_imageset(path=output_dir, processor=detect_resizer())
@cli.command("filehost") @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}") print(f"downloading {meta.suggested_filename} → {output_dir}")
run_command(["aria2c", f"--dir={output_dir}", dl_url]) 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: 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) return any(kw in url for kw in kws)
@cli.command("download") @cli.command("download")
@@ -415,54 +490,10 @@ def download_cmd(
@cli.command("process") @cli.command("process")
@click.argument("paths", nargs=-1, required=True, type=click.Path(path_type=Path)) @click.argument("paths", nargs=-1, required=True, type=click.Path(path_type=Path))
def process_cmd(paths: tuple[Path, ...]) -> None: def process_cmd(paths: tuple[Path, ...]) -> None:
"""Process directories and ZIP files, extracting ZIPs into sibling directories.""" """Process directories and archives, extracting archives into sibling directories."""
for p in paths: processor = detect_resizer()
archive_path: Path | None = None for path in paths:
if not p.exists(): process_imageset(path=path, processor=processor)
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])
if __name__ == "__main__": if __name__ == "__main__":
+217 -2
View File
@@ -3,9 +3,12 @@
# requires-python = ">=3.11" # requires-python = ">=3.11"
# dependencies = ["pytest", "httpx", "pillow", "click"] # dependencies = ["pytest", "httpx", "pillow", "click"]
# /// # ///
import shutil
import subprocess
import sys import sys
import zipfile import zipfile
from pathlib import Path from pathlib import Path
from functools import partial
from unittest.mock import Mock from unittest.mock import Mock
import pytest import pytest
@@ -43,6 +46,47 @@ def test_process_zip_and_directory(
trash.assert_called_once_with(file_paths=[archive_path]) 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"]) @pytest.mark.parametrize("filename", ["missing", "unsupported.txt", "broken.zip"])
def test_process_skips_invalid_inputs( def test_process_skips_invalid_inputs(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, filename: str 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( def test_process_keeps_zip_when_processing_fails(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch tmp_path: Path, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
) -> None: ) -> None:
archive_path = tmp_path / "gallery.zip" archive_path = tmp_path / "gallery.zip"
with zipfile.ZipFile(file=archive_path, mode="w") as archive: 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)] 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() assert archive_path.is_file()
trash.assert_not_called() 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__": if __name__ == "__main__":