Compare commits
4 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| b1f6408da6 | |||
| 55cf88f66e | |||
| 0c5e73eb17 | |||
| 0325467aba |
Binary file not shown.
@@ -1,21 +1,48 @@
|
|||||||
import os
|
import os
|
||||||
|
import re
|
||||||
|
|
||||||
from lyra_worker.library import safe_name
|
from lyra_worker.library import safe_name
|
||||||
|
|
||||||
_AUDIO_EXT = {".flac", ".mp3", ".m4a", ".opus", ".ogg", ".aac", ".wav"}
|
_AUDIO_EXT = {".flac", ".mp3", ".m4a", ".opus", ".ogg", ".aac", ".wav"}
|
||||||
|
|
||||||
|
# Leading "13.", "13 -", "13)" style track-number prefix streamrip bakes into a bonus-track
|
||||||
|
# filename. Requires an explicit delimiter so a real title that merely starts with a number
|
||||||
|
# (e.g. "99 Luftballons") is left alone.
|
||||||
|
_NUM_PREFIX = re.compile(r"^\s*\d+\s*[.):\-]\s*")
|
||||||
|
|
||||||
def plan_track_names(names: list[str], tracklist: tuple[str, ...]) -> list[tuple[str, str, str, int]]:
|
|
||||||
|
def _clean_extra_title(stem: str, artist: str = "") -> str:
|
||||||
|
"""Derive a display title for a track BEYOND the MusicBrainz tracklist (a bonus track that
|
||||||
|
streamrip named itself, e.g. `13. John Mayer - St. Patrick's Day (Album Version)`). Strip a
|
||||||
|
leading track-number prefix, then a leading `<artist> - ` if present, so the example yields
|
||||||
|
`St. Patrick's Day (Album Version)`. Falls back to the stem if nothing strips."""
|
||||||
|
s = _NUM_PREFIX.sub("", stem).strip()
|
||||||
|
if artist:
|
||||||
|
s = re.sub(r"^" + re.escape(artist) + r"\s*-\s*", "", s, flags=re.IGNORECASE).strip()
|
||||||
|
return s or stem
|
||||||
|
|
||||||
|
|
||||||
|
def plan_track_names(
|
||||||
|
names: list[str], tracklist: tuple[str, ...], artist: str = ""
|
||||||
|
) -> list[tuple[str, str, str, int]]:
|
||||||
"""Map audio filenames -> (old_name, new_name, title, track_number).
|
"""Map audio filenames -> (old_name, new_name, title, track_number).
|
||||||
|
|
||||||
Files are sorted; titles come from `tracklist` by position when available,
|
Files are sorted; titles come from `tracklist` by position when available. A file BEYOND
|
||||||
else the filename stem. Pure — no filesystem or tag I/O.
|
the tracklist (a bonus track streamrip named itself) is canonicalized via `_clean_extra_title`
|
||||||
|
so it gets a clean `## Title.ext` name too; with no tracklist at all the on-disk stem is kept.
|
||||||
|
Pure — no filesystem or tag I/O.
|
||||||
"""
|
"""
|
||||||
plan: list[tuple[str, str, str, int]] = []
|
plan: list[tuple[str, str, str, int]] = []
|
||||||
audio = sorted(n for n in names if os.path.splitext(n)[1].lower() in _AUDIO_EXT)
|
audio = sorted(n for n in names if os.path.splitext(n)[1].lower() in _AUDIO_EXT)
|
||||||
for i, name in enumerate(audio):
|
for i, name in enumerate(audio):
|
||||||
ext = os.path.splitext(name)[1].lower()
|
ext = os.path.splitext(name)[1].lower()
|
||||||
title = tracklist[i] if i < len(tracklist) else os.path.splitext(name)[0]
|
stem = os.path.splitext(name)[0]
|
||||||
|
if i < len(tracklist):
|
||||||
|
title = tracklist[i]
|
||||||
|
elif tracklist: # a genuine bonus track past a known tracklist -> canonicalize its name
|
||||||
|
title = _clean_extra_title(stem, artist)
|
||||||
|
else: # no tracklist resolved at all -> keep the on-disk stem (renumbered)
|
||||||
|
title = stem
|
||||||
new_name = f"{i + 1:02d} {safe_name(title)}{ext}"
|
new_name = f"{i + 1:02d} {safe_name(title)}{ext}"
|
||||||
plan.append((name, new_name, title, i + 1))
|
plan.append((name, new_name, title, i + 1))
|
||||||
return plan
|
return plan
|
||||||
@@ -29,7 +56,7 @@ class MutagenTagger:
|
|||||||
import mutagen
|
import mutagen
|
||||||
|
|
||||||
names = os.listdir(album_dir)
|
names = os.listdir(album_dir)
|
||||||
for old_name, new_name, title, number in plan_track_names(names, target.tracklist):
|
for old_name, new_name, title, number in plan_track_names(names, target.tracklist, target.artist):
|
||||||
try:
|
try:
|
||||||
path = os.path.join(album_dir, old_name)
|
path = os.path.join(album_dir, old_name)
|
||||||
audio = mutagen.File(path, easy=True)
|
audio = mutagen.File(path, easy=True)
|
||||||
|
|||||||
@@ -34,6 +34,53 @@ def _run(coro):
|
|||||||
return _loop.run_until_complete(coro)
|
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:
|
def _flatten_audio(dest: str) -> int:
|
||||||
"""streamrip writes into a nested `Artist - Album [...]/` subfolder. Move any audio
|
"""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 up into `dest`, remove the now-empty subfolders, and return the count of audio
|
||||||
@@ -136,6 +183,7 @@ class StreamripClient:
|
|||||||
|
|
||||||
async def _download(self, album_id: str, dest: str, on_progress: Callable[[float], None]) -> None:
|
async def _download(self, album_id: str, dest: str, on_progress: Callable[[float], None]) -> None:
|
||||||
os.makedirs(dest, exist_ok=True)
|
os.makedirs(dest, exist_ok=True)
|
||||||
|
import streamrip.media.track as _track_mod
|
||||||
from streamrip.client import QobuzClient
|
from streamrip.client import QobuzClient
|
||||||
from streamrip.db import Database, Dummy
|
from streamrip.db import Database, Dummy
|
||||||
from streamrip.media import PendingAlbum
|
from streamrip.media import PendingAlbum
|
||||||
@@ -153,7 +201,21 @@ class StreamripClient:
|
|||||||
album = await pending.resolve()
|
album = await pending.resolve()
|
||||||
if album is None:
|
if album is None:
|
||||||
raise RuntimeError(f"could not resolve Qobuz album {album_id}")
|
raise RuntimeError(f"could not resolve Qobuz album {album_id}")
|
||||||
await album.rip()
|
# 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)
|
on_progress(1.0)
|
||||||
finally:
|
finally:
|
||||||
await client.session.close()
|
await client.session.close()
|
||||||
|
|||||||
@@ -2,6 +2,7 @@ import os
|
|||||||
import shutil
|
import shutil
|
||||||
import threading
|
import threading
|
||||||
import time
|
import time
|
||||||
|
from concurrent.futures import ThreadPoolExecutor
|
||||||
from dataclasses import replace
|
from dataclasses import replace
|
||||||
from typing import Callable, Sequence
|
from typing import Callable, Sequence
|
||||||
|
|
||||||
@@ -94,6 +95,16 @@ def _poll_download_progress(job_id: str, staging: str, expected: int, stop: "thr
|
|||||||
conn.close()
|
conn.close()
|
||||||
|
|
||||||
|
|
||||||
|
def _search_adapter(adapter: SourceAdapter, target: MBTarget) -> list[Candidate]:
|
||||||
|
"""Search one adapter, swallowing failures — a down source contributes no candidates rather
|
||||||
|
than crashing the job. Runs on a worker thread (searches are parallelized across sources)."""
|
||||||
|
try:
|
||||||
|
return list(adapter.search(target))
|
||||||
|
except Exception as e:
|
||||||
|
print(f"pipeline: adapter {adapter.name} search failed: {e}", flush=True)
|
||||||
|
return []
|
||||||
|
|
||||||
|
|
||||||
def _set_state(conn: psycopg.Connection, job_id: str, state: str, stage: str) -> None:
|
def _set_state(conn: psycopg.Connection, job_id: str, state: str, stage: str) -> None:
|
||||||
with conn.cursor() as cur:
|
with conn.cursor() as cur:
|
||||||
cur.execute(
|
cur.execute(
|
||||||
@@ -306,15 +317,16 @@ def run_pipeline(
|
|||||||
conn.commit()
|
conn.commit()
|
||||||
return
|
return
|
||||||
|
|
||||||
# 2. match — search all adapters; score + persist EVERY candidate found
|
# 2. match — search all adapters CONCURRENTLY (Soulseek's search polls up to ~2 min; run it
|
||||||
|
# alongside Qobuz's instant search so total match time = the slowest source, not their sum).
|
||||||
|
# Only Qobuz touches streamrip's shared asyncio loop and there's a single Qobuz adapter, so no
|
||||||
|
# two threads drive that loop at once. Results are collected in adapter order (map preserves
|
||||||
|
# input order) so candidate ordering stays deterministic.
|
||||||
_set_state(conn, job_id, "matching", "match")
|
_set_state(conn, job_id, "matching", "match")
|
||||||
found: list[Candidate] = []
|
found: list[Candidate] = []
|
||||||
for adapter in adapters:
|
with ThreadPoolExecutor(max_workers=max(1, len(adapters))) as pool:
|
||||||
try:
|
per_adapter = list(pool.map(lambda a: (a, _search_adapter(a, target)), adapters))
|
||||||
results = adapter.search(target)
|
for adapter, results in per_adapter:
|
||||||
except Exception as e: # a down source contributes no candidates, never crashes the job
|
|
||||||
print(f"pipeline: adapter {adapter.name} search failed: {e}", flush=True)
|
|
||||||
continue
|
|
||||||
for c in results:
|
for c in results:
|
||||||
found.append(replace(c, source_tier=adapter.tier, confidence=score_confidence(target, c)))
|
found.append(replace(c, source_tier=adapter.tier, confidence=score_confidence(target, c)))
|
||||||
_persist_candidates(conn, job_id, found)
|
_persist_candidates(conn, job_id, found)
|
||||||
|
|||||||
@@ -1,4 +1,34 @@
|
|||||||
from lyra_worker._mutagen import plan_track_names
|
from lyra_worker._mutagen import _clean_extra_title, plan_track_names
|
||||||
|
|
||||||
|
|
||||||
|
def test_clean_extra_title_strips_number_and_artist_prefix():
|
||||||
|
assert _clean_extra_title(
|
||||||
|
"13. John Mayer - St. Patrick's Day (Album Version)", "John Mayer"
|
||||||
|
) == "St. Patrick's Day (Album Version)"
|
||||||
|
|
||||||
|
|
||||||
|
def test_clean_extra_title_strips_number_without_artist_prefix():
|
||||||
|
assert _clean_extra_title("13. St. Patrick's Day", "John Mayer") == "St. Patrick's Day"
|
||||||
|
|
||||||
|
|
||||||
|
def test_clean_extra_title_leaves_plain_title_and_numeric_titles():
|
||||||
|
assert _clean_extra_title("Hidden Track", "John Mayer") == "Hidden Track"
|
||||||
|
# "99 Luftballons" has no delimiter after the number -> not a track-number prefix
|
||||||
|
assert _clean_extra_title("99 Luftballons", "Nena") == "99 Luftballons"
|
||||||
|
|
||||||
|
|
||||||
|
def test_plan_canonicalizes_bonus_track_beyond_tracklist():
|
||||||
|
names = [f"{i:02d} Track{i}.flac" for i in range(1, 13)] + [
|
||||||
|
"13. John Mayer - St. Patrick's Day (Album Version).flac"
|
||||||
|
]
|
||||||
|
tracklist = tuple(f"Track{i}" for i in range(1, 13))
|
||||||
|
plan = plan_track_names(names, tracklist, artist="John Mayer")
|
||||||
|
assert plan[12] == (
|
||||||
|
"13. John Mayer - St. Patrick's Day (Album Version).flac",
|
||||||
|
"13 St. Patrick's Day (Album Version).flac",
|
||||||
|
"St. Patrick's Day (Album Version)",
|
||||||
|
13,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def test_plan_uses_tracklist_titles_and_numbers():
|
def test_plan_uses_tracklist_titles_and_numbers():
|
||||||
|
|||||||
@@ -195,6 +195,33 @@ def test_below_threshold_candidates_are_persisted_for_diagnosis(conn):
|
|||||||
assert cur.fetchone()[0] >= 1 # found-but-rejected candidate retained for diagnosis
|
assert cur.fetchone()[0] >= 1 # found-but-rejected candidate retained for diagnosis
|
||||||
|
|
||||||
|
|
||||||
|
def test_searches_run_concurrently_across_adapters(conn):
|
||||||
|
# Two sources that each take ~0.5s to search must run in parallel: total match time is ~the
|
||||||
|
# slowest single search (~0.5s), not the sum (~1.0s). Also verifies concurrency drops no
|
||||||
|
# candidates — both sources' contributions are persisted.
|
||||||
|
import time as _t
|
||||||
|
|
||||||
|
class _SlowQobuz(FakeQobuz):
|
||||||
|
def search(self, target):
|
||||||
|
_t.sleep(0.5)
|
||||||
|
return super().search(target)
|
||||||
|
|
||||||
|
class _SlowSoulseek(FakeSoulseek):
|
||||||
|
def search(self, target):
|
||||||
|
_t.sleep(0.5)
|
||||||
|
return super().search(target)
|
||||||
|
|
||||||
|
job_id = insert_request(conn, artist="Radiohead", album="In Rainbows")
|
||||||
|
claim_next(conn)
|
||||||
|
start = _t.monotonic()
|
||||||
|
run_pipeline(conn, job_id, [_SlowQobuz(), _SlowSoulseek()], dest_root="/tmp/lib")
|
||||||
|
elapsed = _t.monotonic() - start
|
||||||
|
assert elapsed < 0.9 # concurrent (~0.5s), not sequential (~1.0s)
|
||||||
|
with conn.cursor() as cur:
|
||||||
|
cur.execute('SELECT count(DISTINCT source) FROM "Candidate" WHERE "jobId" = %s', (job_id,))
|
||||||
|
assert cur.fetchone()[0] == 2 # both adapters contributed candidates
|
||||||
|
|
||||||
|
|
||||||
def test_duplicate_adapter_names_raise(conn):
|
def test_duplicate_adapter_names_raise(conn):
|
||||||
import pytest
|
import pytest
|
||||||
from lyra_worker.adapters.fakes import FakeQobuz
|
from lyra_worker.adapters.fakes import FakeQobuz
|
||||||
|
|||||||
@@ -0,0 +1,41 @@
|
|||||||
|
from lyra_worker.adapters._streamrip import _ProgressAggregator
|
||||||
|
|
||||||
|
|
||||||
|
def _record():
|
||||||
|
seen: list[float] = []
|
||||||
|
return seen, seen.append
|
||||||
|
|
||||||
|
|
||||||
|
def test_equal_size_tracks_aggregate_monotonically_to_near_one():
|
||||||
|
# 4 tracks of 100 bytes each, all started up front (concurrency covers the album), then
|
||||||
|
# each downloaded in 25-byte chunks. Fraction must climb monotonically 0→~1.
|
||||||
|
seen, cb = _record()
|
||||||
|
agg = _ProgressAggregator(4, cb)
|
||||||
|
for _ in range(4):
|
||||||
|
agg.start_track(100)
|
||||||
|
for _ in range(16): # 16 chunks * 25 bytes = 400 bytes = full album
|
||||||
|
agg.advance(25)
|
||||||
|
|
||||||
|
assert seen == sorted(seen) # monotonic non-decreasing
|
||||||
|
assert all(0.0 <= f <= 0.99 for f in seen) # never negative, capped at 0.99
|
||||||
|
assert seen[-1] == 0.99 # reaches the cap when all bytes are in
|
||||||
|
assert seen[len(seen) // 2] > 0.3 # genuinely climbs mid-download (no 0→100 jump)
|
||||||
|
|
||||||
|
|
||||||
|
def test_estimate_uses_track_count_before_all_tracks_start():
|
||||||
|
# Only the first of 4 tracks has started; with a known track count the estimate already
|
||||||
|
# spans the whole album, so finishing track 1 reports ~1/4, not ~1/1.
|
||||||
|
seen, cb = _record()
|
||||||
|
agg = _ProgressAggregator(4, cb)
|
||||||
|
agg.start_track(100)
|
||||||
|
agg.advance(100) # track 1 fully downloaded
|
||||||
|
assert abs(seen[-1] - 0.25) < 0.01
|
||||||
|
|
||||||
|
|
||||||
|
def test_zero_or_missing_sizes_never_emit_and_never_divide_by_zero():
|
||||||
|
seen, cb = _record()
|
||||||
|
agg = _ProgressAggregator(0, cb)
|
||||||
|
agg.advance(50) # no track started yet -> no emit
|
||||||
|
agg.start_track(0) # a track with unknown size
|
||||||
|
agg.advance(50) # est_total still 0 -> still no emit, no crash
|
||||||
|
assert seen == []
|
||||||
Reference in New Issue
Block a user