feat(download_imageset): Use tricklebar for download when possible

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())
 

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,