Jonathan
0325467aba
streamrip reports progress only per concurrent track, so Job.downloadProgress came solely from the file-count poller — 0% until every track landed, then a jump to 100%. Patch streamrip.media.track.get_progress_callback to aggregate per-track byte deltas over an estimated album total (avg started-track size × len(album.tracks)) into one smooth, generally-climbing 0..1 fraction routed through the pipeline's on_progress. Capped at 0.99; the pipeline sets 1.0 on completion. No signature changes — track count read from the resolved album. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
222 lines
9.8 KiB
Python
222 lines
9.8 KiB
Python
import asyncio
|
||
import hashlib
|
||
import http.client
|
||
import os
|
||
from typing import Callable
|
||
|
||
# Qobuz's CDN returns responses with more than Python's default 100-header limit for some
|
||
# tracks, so http.client raises HTTPException("got more than 100 headers"), streamrip skips
|
||
# that track, and it lands as a 0-byte file (seen on Bruno Mars "24K Magic" — track 1). Raise
|
||
# the limit process-wide so those downloads succeed. (Set on import of the Qobuz adapter.)
|
||
http.client._MAXHEADERS = 1000
|
||
|
||
|
||
def _hashed_password(password: str) -> str:
|
||
"""streamrip's Qobuz email/password login expects the password as an MD5 hex digest."""
|
||
return hashlib.md5(password.encode("utf-8")).hexdigest()
|
||
|
||
|
||
_AUDIO_EXT = {".flac", ".mp3", ".m4a", ".opus", ".ogg", ".aac", ".wav"}
|
||
|
||
# streamrip holds module-global asyncio primitives (a download-concurrency Semaphore) that
|
||
# bind to the running event loop on first use. `asyncio.run` creates a NEW loop per call and
|
||
# closes it, orphaning that Semaphore, so the 2nd+ download fails with "Semaphore ... is bound
|
||
# to a different event loop" and its tracks silently fail (-> incomplete download). Running every
|
||
# streamrip coroutine on ONE persistent loop keeps those globals valid across downloads. The
|
||
# worker processes one job at a time, so a single shared loop is safe.
|
||
_loop: "asyncio.AbstractEventLoop | None" = None
|
||
|
||
|
||
def _run(coro):
|
||
global _loop
|
||
if _loop is None or _loop.is_closed():
|
||
_loop = asyncio.new_event_loop()
|
||
return _loop.run_until_complete(coro)
|
||
|
||
|
||
class _ProgressAggregator:
|
||
"""Turn streamrip's per-track byte callbacks into one album-level 0..1 fraction.
|
||
|
||
streamrip downloads an album's tracks concurrently and reports progress only per track
|
||
(a byte delta per chunk), so the file-count poller sees 0 complete files until they all
|
||
land at once — the 0%→100% jump. Summing downloaded bytes over an *estimated* album total
|
||
(average started-track size × track count) gives a smooth, generally-climbing fraction.
|
||
Track sizes are similar, so the estimate is stable once a track or two has started; capped
|
||
at 0.99 because the pipeline sets 1.0 once the whole download returns. All streamrip
|
||
coroutines run on one event loop on the calling thread, so no locking is needed."""
|
||
|
||
def __init__(self, track_count: int, on_progress: Callable[[float], None]):
|
||
self._n = track_count or 0
|
||
self._on_progress = on_progress
|
||
self._seen_total = 0 # summed byte-size of tracks that have started downloading
|
||
self._seen_count = 0 # number of tracks that have started
|
||
self._done_bytes = 0 # summed bytes downloaded across all started tracks
|
||
|
||
def start_track(self, total) -> None:
|
||
self._seen_total += max(int(total or 0), 0)
|
||
self._seen_count += 1
|
||
|
||
def advance(self, delta) -> None:
|
||
self._done_bytes += max(int(delta or 0), 0)
|
||
if self._seen_count == 0:
|
||
return
|
||
avg = self._seen_total / self._seen_count
|
||
n = self._n if self._n > 0 else self._seen_count
|
||
est_total = avg * n
|
||
if est_total > 0:
|
||
self._on_progress(min(self._done_bytes / est_total, 0.99))
|
||
|
||
|
||
class _ProgressHandle:
|
||
"""streamrip's `Track.download` uses `with get_progress_callback(...) as cb:`; `cb` must be
|
||
the per-chunk byte callback. This stand-in yields that callback and no-ops on exit."""
|
||
|
||
def __init__(self, update: Callable[[int], None]):
|
||
self.update = update
|
||
|
||
def __enter__(self):
|
||
return self.update
|
||
|
||
def __exit__(self, *_):
|
||
return False
|
||
|
||
|
||
def _flatten_audio(dest: str) -> int:
|
||
"""streamrip writes into a nested `Artist - Album [...]/` subfolder. Move any audio
|
||
files up into `dest`, remove the now-empty subfolders, and return the count of audio
|
||
files directly in `dest` — the REAL number downloaded, so a partial download (streamrip
|
||
continues past per-track failures) reports its true short count to the completeness check.
|
||
"""
|
||
dest = os.path.abspath(dest)
|
||
for root, _dirs, files in os.walk(dest):
|
||
if os.path.abspath(root) == dest:
|
||
continue
|
||
for name in files:
|
||
if os.path.splitext(name)[1].lower() not in _AUDIO_EXT:
|
||
continue
|
||
src = os.path.join(root, name)
|
||
target = os.path.join(dest, name)
|
||
if os.path.exists(target):
|
||
base, ext = os.path.splitext(name)
|
||
target = os.path.join(dest, f"{base} ({os.path.basename(root)}){ext}")
|
||
os.rename(src, target)
|
||
for root, _dirs, _files in os.walk(dest, topdown=False):
|
||
if os.path.abspath(root) != dest:
|
||
try:
|
||
os.rmdir(root)
|
||
except OSError:
|
||
pass # subfolder still has non-audio files (e.g. cover art) — leave it
|
||
return sum(1 for f in os.listdir(dest) if os.path.splitext(f)[1].lower() in _AUDIO_EXT)
|
||
|
||
|
||
class StreamripClient:
|
||
"""Real Qobuz client backed by streamrip (v2, async → sync via asyncio.run).
|
||
|
||
NOT unit-tested offline; see test_qobuz_live.py. streamrip is imported lazily
|
||
inside the async helpers so importing/constructing this class stays offline.
|
||
"""
|
||
|
||
def __init__(self, config: dict):
|
||
self._email = config.get("qobuz.email", "")
|
||
self._password = config.get("qobuz.password", "")
|
||
self._user_id = config.get("qobuz.user_id", "")
|
||
self._token = config.get("qobuz.auth_token", "")
|
||
|
||
def is_configured(self) -> bool:
|
||
# either a user_id + auth token, or an email + password
|
||
return bool((self._user_id and self._token) or (self._email and self._password))
|
||
|
||
def _make_config(self, download_folder: str | None = None):
|
||
from streamrip.config import Config
|
||
|
||
cfg = Config.defaults()
|
||
q = cfg.session.qobuz
|
||
if self._user_id and self._token:
|
||
# token auth: more reliable than email/password, which Qobuz often rejects
|
||
q.use_auth_token = True
|
||
q.email_or_userid = self._user_id
|
||
q.password_or_token = self._token
|
||
else:
|
||
q.use_auth_token = False
|
||
q.email_or_userid = self._email
|
||
q.password_or_token = _hashed_password(self._password)
|
||
q.quality = 3
|
||
if download_folder is not None:
|
||
cfg.session.downloads.folder = download_folder
|
||
return cfg
|
||
|
||
def search_album(self, artist: str, album: str) -> list[dict]:
|
||
return _run(self._search(f"{artist} {album}"))
|
||
|
||
async def _search(self, query: str) -> list[dict]:
|
||
from streamrip.client import QobuzClient
|
||
|
||
client = QobuzClient(self._make_config())
|
||
try:
|
||
try:
|
||
await client.login()
|
||
except Exception as e: # never let streamrip's credential-bearing error escape
|
||
raise RuntimeError(f"Qobuz login failed ({type(e).__name__}) — check credentials") from None
|
||
pages = await client.search("album", query, limit=5)
|
||
finally:
|
||
await client.session.close()
|
||
results: list[dict] = []
|
||
for page in pages:
|
||
for item in page.get("albums", {}).get("items", []):
|
||
results.append(
|
||
{
|
||
"source_ref": str(item.get("id", "")),
|
||
"title": item.get("title", ""),
|
||
"artist": (item.get("artist") or {}).get("name", ""),
|
||
"track_count": item.get("tracks_count", 0) or 0,
|
||
"bit_depth": item.get("maximum_bit_depth", 16) or 16,
|
||
"sampling_rate": item.get("maximum_sampling_rate", 44.1) or 44.1,
|
||
}
|
||
)
|
||
return results
|
||
|
||
def download(self, source_ref: str, dest: str, on_progress: Callable[[float], None]) -> dict:
|
||
_run(self._download(source_ref, dest, on_progress))
|
||
# Count the files actually written and flatten streamrip's nested output into dest
|
||
# so the tag stage can organize it. A partial download -> short count -> needs_attention.
|
||
return {"track_count": _flatten_audio(dest), "path": dest}
|
||
|
||
async def _download(self, album_id: str, dest: str, on_progress: Callable[[float], None]) -> None:
|
||
os.makedirs(dest, exist_ok=True)
|
||
import streamrip.media.track as _track_mod
|
||
from streamrip.client import QobuzClient
|
||
from streamrip.db import Database, Dummy
|
||
from streamrip.media import PendingAlbum
|
||
|
||
cfg = self._make_config(download_folder=dest)
|
||
client = QobuzClient(cfg)
|
||
try:
|
||
try:
|
||
await client.login()
|
||
except Exception as e: # never let streamrip's credential-bearing error escape
|
||
raise RuntimeError(f"Qobuz login failed ({type(e).__name__}) — check credentials") from None
|
||
on_progress(0.0)
|
||
db = Database(downloads=Dummy(), failed=Dummy())
|
||
pending = PendingAlbum(album_id, client, cfg, db)
|
||
album = await pending.resolve()
|
||
if album is None:
|
||
raise RuntimeError(f"could not resolve Qobuz album {album_id}")
|
||
# Hook streamrip's per-track byte progress into a single album-level fraction so the
|
||
# Floor bar climbs smoothly instead of jumping 0→100. `Track.download` binds
|
||
# get_progress_callback at import, so patch it on the track module (not progress.py).
|
||
agg = _ProgressAggregator(len(getattr(album, "tracks", []) or []), on_progress)
|
||
_orig_cb = _track_mod.get_progress_callback
|
||
|
||
def _hook(_enabled, total, _desc):
|
||
agg.start_track(total)
|
||
return _ProgressHandle(agg.advance)
|
||
|
||
_track_mod.get_progress_callback = _hook
|
||
try:
|
||
await album.rip()
|
||
finally:
|
||
_track_mod.get_progress_callback = _orig_cb
|
||
on_progress(1.0)
|
||
finally:
|
||
await client.session.close()
|