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