diff --git a/download_imageset.py b/download_imageset.py index 1f75c51..6f7d490 100755 --- a/download_imageset.py +++ b/download_imageset.py @@ -7,9 +7,11 @@ import dataclasses import logging import os import re +import shlex import shutil import struct import subprocess +import tempfile import urllib.parse import zipfile from functools import partial @@ -201,6 +203,62 @@ def _has_images(d: Path) -> bool: return any(f.suffix.lower() in IMAGE_EXTS for f in d.rglob("*") if f.is_file()) +_TRICKLEBAR_APP = Path("/Applications/TrickleBar.app") + + +def download_url(url: str, destination: Path) -> None: + destination = destination.resolve() + destination.parent.mkdir(parents=True, exist_ok=True) + if _TRICKLEBAR_APP.is_dir(): + logging.info(f"downloading with TrickleBar → {destination}") + download_url_with_tricklebar(url=url, destination=destination) + else: + logging.info(f"downloading with aria2 → {destination}") + download_url_with_aria2(url=url, destination=destination) + + +def download_url_with_aria2(url: str, destination: Path) -> None: + run_command(args=[ + "aria2c", "--continue=true", f"--dir={destination.parent}", + f"--out={destination.name}", url, + ]) + + +def download_url_with_tricklebar(url: str, destination: Path) -> None: + with tempfile.TemporaryDirectory(prefix="imageset-") as temporary_directory: + directory = Path(temporary_directory) + completion_pipe = directory / "completion" + callback = directory / "on_done.sh" + os.mkfifo(path=completion_pipe, mode=0o600) + callback.write_text(data=( + "#!/bin/sh\n" + f"[ -p {shlex.quote(str(completion_pipe))} ] || exit 0\n" + f"printf '%s\\0' \"$1\" \"$2\" > {shlex.quote(str(completion_pipe))}\n" + )) + callback.chmod(mode=0o700) + query = urllib.parse.urlencode(query={ + "on_done": str(callback), "url": url, + "dir": str(destination.parent), "out": destination.name, "continue": "true", + }, quote_via=urllib.parse.quote) + logging.info("waiting for TrickleBar; press Ctrl+C to stop waiting") + run_command(args=["open", f"tricklebar://add-download?{query}"]) + with completion_pipe.open(mode="rb") as completion: + status, downloaded_path, _ = completion.read().split(b"\0", maxsplit=2) + if status != b"complete": + raise click.ClickException( + f"TrickleBar download failed: {os.fsdecode(status)}" + ) + if not downloaded_path or not Path(os.fsdecode(downloaded_path)).is_file(): + raise click.ClickException("TrickleBar completed without a downloaded file") + if Path(os.fsdecode(downloaded_path)).resolve() != destination.resolve(): + raise click.ClickException( + f"TrickleBar downloaded to {os.fsdecode(downloaded_path)} instead of {destination}; " + "update TrickleBar to a version supporting dir/out URL options" + ) + if destination.with_name(f"{destination.name}.aria2").exists(): + raise click.ClickException(f"TrickleBar download is unfinished: {destination}") + + def _needs_resize(d: Path) -> bool: 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: @@ -463,10 +521,7 @@ def filehost_cmd(urls: tuple[str, ...], name: str | None, cwd: Path, token: str) ): logging.info(f"skipping: {video_path.name} already downloaded") continue - run_command(args=[ - "aria2c", "--continue=true", f"--dir={video_path.parent}", - f"--out={video_path.name}", dl_url, - ]) + download_url(url=dl_url, destination=video_path) continue dest = output_dir / meta.suggested_filename @@ -480,10 +535,7 @@ def filehost_cmd(urls: tuple[str, ...], name: str | None, cwd: Path, token: str) logging.info(f"skipping: {dest.name} already downloaded") continue logging.info(f"downloading {meta.suggested_filename} → {output_dir}") - run_command(args=[ - "aria2c", "--continue=true", f"--dir={output_dir}", - f"--out={dest.name}", dl_url, - ]) + download_url(url=dl_url, destination=dest) process_imageset(path=output_dir, processor=detect_resizer()) diff --git a/download_imageset_test.py b/download_imageset_test.py index 58cace0..66645ca 100755 --- a/download_imageset_test.py +++ b/download_imageset_test.py @@ -6,7 +6,10 @@ import shutil import subprocess import sys +import threading +import urllib.parse import zipfile +from concurrent.futures import ThreadPoolExecutor from pathlib import Path from functools import partial from unittest.mock import Mock @@ -330,6 +333,7 @@ def test_filehost_resumes_unfinished_downloads( )), ) run_command = Mock() + monkeypatch.setattr(download_imageset, "_TRICKLEBAR_APP", tmp_path / "missing.app") monkeypatch.setattr(download_imageset, "run_command", run_command) monkeypatch.setattr(download_imageset, "process_imageset", Mock()) @@ -350,6 +354,97 @@ def test_filehost_resumes_unfinished_downloads( run_command.assert_not_called() +def test_aria2_download_failure_propagates(monkeypatch: pytest.MonkeyPatch) -> None: + failure = subprocess.CalledProcessError(returncode=1, cmd="aria2c") + monkeypatch.setattr(download_imageset, "run_command", Mock(side_effect=failure)) + + with pytest.raises(subprocess.CalledProcessError): + download_imageset.download_url_with_aria2( + url="https://example.com/archive.zip", destination=Path("archive.zip") + ) + + +@pytest.mark.parametrize( + ("status", "file_state", "error"), + [ + ("complete", "valid", None), + ("error", "valid", "download failed"), + ("complete", "missing", "without a downloaded file"), + ("complete", "wrong_path", "instead of"), + ("complete", "unfinished", "download is unfinished"), + ], +) +def test_tricklebar_waits_for_callback( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, + status: str, file_state: str, error: str | None, +) -> None: + destination = tmp_path / "download & photos.zip" + downloaded_path = tmp_path / "wrong.zip" if file_state == "wrong_path" else destination + if file_state != "missing": + downloaded_path.write_bytes(data=b"downloaded archive") + if file_state == "unfinished": + destination.with_name(f"{destination.name}.aria2").touch() + url = "https://example.com/photos.zip?token=a+b&other=value" + opened = threading.Event() + callback_paths: list[Path] = [] + + def open_download(args: list[str]) -> None: + assert args[0] == "open" + parsed = urllib.parse.urlparse(url=args[1]) + assert "+" not in parsed.query + assert parsed.scheme == "tricklebar" + assert parsed.netloc == "add-download" + query = urllib.parse.parse_qs(qs=parsed.query) + assert query["url"] == [url] + assert query["dir"] == [str(destination.parent)] + assert query["out"] == [destination.name] + assert query["continue"] == ["true"] + callback_paths.append(Path(query["on_done"][0])) + opened.set() + + monkeypatch.setattr(download_imageset, "run_command", open_download) + + with ThreadPoolExecutor(max_workers=1) as executor: + download = executor.submit( + download_imageset.download_url_with_tricklebar, url=url, destination=destination + ) + assert opened.wait(timeout=5) + assert not download.done() + callback = callback_paths[0] + subprocess.run(args=[str(callback), status, str(downloaded_path), "gid", url], check=True) + if error is None: + assert download.result(timeout=5) is None + else: + with pytest.raises(download_imageset.click.ClickException, match=error): + download.result(timeout=5) + + assert not callback.parent.exists() + + +@pytest.mark.parametrize("installed", [False, True]) +def test_download_selects_backend( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, installed: bool +) -> None: + application = tmp_path / "TrickleBar.app" + if installed: + application.mkdir() + monkeypatch.setattr(download_imageset, "_TRICKLEBAR_APP", application) + aria2 = Mock() + tricklebar = Mock() + monkeypatch.setattr(download_imageset, "download_url_with_aria2", aria2) + monkeypatch.setattr(download_imageset, "download_url_with_tricklebar", tricklebar) + destination = tmp_path / "photos" / "album.zip" + url = "https://example.com/album.zip" + + download_imageset.download_url(url=url, destination=destination) + + assert destination.parent.is_dir() + selected = tricklebar if installed else aria2 + unused = aria2 if installed else tricklebar + selected.assert_called_once_with(url=url, destination=destination.resolve()) + unused.assert_not_called() + + @pytest.mark.parametrize("process_directory", [False, True]) def test_process_skips_unfinished_archive( tmp_path: Path,