diff --git a/core/upgrader.py b/core/upgrader.py index 822ca9d..da8c4b4 100644 --- a/core/upgrader.py +++ b/core/upgrader.py @@ -5,15 +5,22 @@ from __future__ import annotations import json import queue import re +import shutil import subprocess import threading from pathlib import Path from typing import Callable, Optional -from core.paths import find_ytdlp, find_ffmpeg_dir, find_ffprobe, subprocess_flags +from core.paths import find_ytdlp, find_ffmpeg_dir, find_ffmpeg, find_ffprobe, subprocess_flags AUDIO_EXTENSIONS = {".mp3", ".m4a", ".wav", ".flac"} +# Soglia minima Jaccard per considerare un file dell'archivio come candidato. +_ARCHIVE_MIN_SIMILARITY = 0.5 + +# Token di lunghezza inferiore a questa vengono scartati (troppo generici). +_MIN_TOKEN_LEN = 3 + # Timeout massimo per singolo download yt-dlp (secondi). Watchdog kill. # Serve a evitare hang su video geo-restricted, YouTube throttle o rete lenta. _DOWNLOAD_TIMEOUT_SEC = 300 @@ -75,6 +82,140 @@ def get_bitrate(filepath: str | Path) -> int: return 0 +def _normalize_stem(stem: str) -> set: + """Ritorna il set di token normalizzati per un nome file (senza extension). + + Pipeline: + 1. Lowercase + 2. Sostituisce caratteri non alfanumerici con spazi (mantiene split naturale) + 3. Split su whitespace + 4. Filtra token di lunghezza < _MIN_TOKEN_LEN (troppo generici) + """ + if not stem: + return set() + lowered = stem.lower() + # Manteniamo spazi ma sostituiamo tutto il resto (che non è alfanumerico) con spazi + cleaned = re.sub(r"[^a-z0-9\s]+", " ", lowered) + tokens = cleaned.split() + return {t for t in tokens if len(t) >= _MIN_TOKEN_LEN} + + +def _tokens_to_key(tokens: set) -> str: + """Chiave stabile per un set di token (usata come chiave del dict d'indice).""" + return "|".join(sorted(tokens)) + + +def _key_to_tokens(key: str) -> set: + """Inverso di _tokens_to_key.""" + if not key: + return set() + return set(key.split("|")) + + +def _scan_archive(archive_dir: str) -> dict: + """Indicizza recursivamente un archivio di file audio. + + Ritorna dict {token_key: [Path, ...]}. La scansione è case-insensitive + sulle estensioni: `.MP3`, `.Mp3`, `.mp3` sono tutti riconosciuti. + """ + base = Path(archive_dir) + index: dict = {} + if not base.exists() or not base.is_dir(): + return index + for f in base.rglob("*"): + if not f.is_file(): + continue + if f.suffix.lower() not in AUDIO_EXTENSIONS: + continue + tokens = _normalize_stem(f.stem) + if not tokens: + continue + key = _tokens_to_key(tokens) + index.setdefault(key, []).append(f) + return index + + +def _find_candidates( + target_stem: str, + index: dict, + min_similarity: float = _ARCHIVE_MIN_SIMILARITY, +) -> list: + """Cerca candidati nell'indice tramite Jaccard sui token del nome. + + Ritorna lista di tuple `(path, similarity, bitrate)` ordinata per + bitrate DESC, poi similarity DESC. + """ + target_tokens = _normalize_stem(target_stem) + if not target_tokens: + return [] + + results: list = [] + for key, paths in index.items(): + entry_tokens = _key_to_tokens(key) + if not entry_tokens: + continue + common = len(target_tokens & entry_tokens) + if common == 0: + continue + total = len(target_tokens | entry_tokens) + if total == 0: + continue + sim = common / total + if sim < min_similarity: + continue + for p in paths: + try: + br = get_bitrate(p) + except Exception: + br = 0 + results.append((p, sim, br)) + + # Ordina: bitrate DESC, poi similarity DESC (stabile) + results.sort(key=lambda t: (-t[2], -t[1])) + return results + + +def _copy_or_convert_to_mp3( + src: Path, + dst: Path, + temp_dir: Path, +) -> bool: + """Copia (o converte se necessario) `src` in `dst` come MP3. + + - Se src è già .mp3 → copia diretta con shutil.copy2 + - Altrimenti → converti con ffmpeg a 320k CBR + Ritorna True in caso di successo. + """ + try: + if src.suffix.lower() == ".mp3": + shutil.copy2(str(src), str(dst)) + return dst.exists() + # Convert non-mp3 -> mp3 320k + ffmpeg_bin = find_ffmpeg() or "ffmpeg" + temp_out = temp_dir / f"_archive_convert_{dst.stem}.mp3" + # -y per sovrascrivere se residuo di run precedente + subprocess.run( + [ + ffmpeg_bin, "-y", + "-i", str(src), + "-vn", # scarta eventuali stream video/cover art (le riproviamo dopo) + "-c:a", "libmp3lame", + "-b:a", "320k", + "-id3v2_version", "3", + str(temp_out), + ], + capture_output=True, + timeout=300, + **subprocess_flags(), + ) + if not temp_out.exists(): + return False + temp_out.replace(dst) + return True + except Exception: + return False + + def _load_done_set(done_file: Path) -> set[str]: if done_file.exists(): return set(done_file.read_text(encoding="utf-8").splitlines()) @@ -185,8 +326,21 @@ def upgrade_folder( cookies_path: Optional[str] = None, recursive: bool = False, progress_callback: Optional[Callable] = None, + archive_dir: Optional[str] = None, + resolve_callback: Optional[Callable] = None, ) -> None: - """Logica principale di upgrade qualita.""" + """Logica principale di upgrade qualita. + + Args: + archive_dir: se presente, prima di scaricare da YouTube l'app cerca + una versione HQ del brano in questa cartella (recursive). Se ne + trova una la copia (o converte in mp3 320k) mantenendo il nome + originale. + resolve_callback: chiamato in caso di match multiplo nell'archivio. + Riceve una lista di dict {path, bitrate, size, similarity} e deve + ritornare bloccante un dict {'action': 'use_local'|'use_youtube'| + 'skip', 'path': Optional[str]}. + """ reset_stop() ytdlp = find_ytdlp() @@ -200,7 +354,7 @@ def upgrade_folder( else: folders = [Path(directory)] - all_items: list[tuple[Path, Path]] = [] + all_items: list = [] for folder in folders: for ext in AUDIO_EXTENSIONS: for f in sorted(folder.glob(f"*{ext}")): @@ -212,6 +366,16 @@ def upgrade_folder( progress_callback(0, 0, "", "no_files", 0, 0) return + # Pre-scan dell'archivio (una volta sola). Se archive_dir è None si salta. + archive_index: dict = {} + if archive_dir: + try: + archive_index = _scan_archive(archive_dir) + except Exception: + archive_index = {} + if progress_callback: + progress_callback(0, total, "", "scan_archive", 0, len(archive_index)) + processed = 0 for filepath, folder in all_items: @@ -235,6 +399,95 @@ def upgrade_folder( current_kbps = get_bitrate(filepath) + # ------------------------------------------------------------------ + # 1) ARCHIVE LOOKUP (se archive_dir presente) + # ------------------------------------------------------------------ + if archive_index: + try: + candidates = _find_candidates(filename, archive_index) + except Exception: + candidates = [] + + selected_path: Optional[Path] = None + user_chose_youtube = False + user_chose_skip = False + + if len(candidates) == 1: + selected_path = candidates[0][0] + elif len(candidates) >= 2 and resolve_callback is not None: + cand_payload = [] + for cp, csim, cbr in candidates: + try: + csize = cp.stat().st_size + except Exception: + csize = 0 + cand_payload.append({ + "path": str(cp), + "bitrate": cbr, + "size": csize, + "similarity": csim, + }) + if progress_callback: + progress_callback(processed, total, filepath.name, + "resolve_wait", current_kbps, 0) + try: + choice = resolve_callback(filepath.name, cand_payload) or {} + except Exception: + choice = {} + action = (choice.get("action") or "").strip() + if action == "use_local": + chosen = (choice.get("path") or "").strip() + if chosen: + cp = Path(chosen) + if cp.exists(): + selected_path = cp + elif action == "use_youtube": + user_chose_youtube = True + elif action == "skip": + user_chose_skip = True + else: + # Risposta invalida: fallback su YouTube per non bloccare + user_chose_youtube = True + # len(candidates) >= 2 senza callback: fallback YouTube + # len(candidates) == 0: fallback YouTube + + if user_chose_skip: + _mark_done(done_file, filename) + processed += 1 + if progress_callback: + progress_callback(processed, total, filepath.name, + "skipped_by_user", current_kbps, 0) + continue + + if selected_path is not None and not user_chose_youtube: + # Copia/converte il candidato locale come .mp3 con nome originale + dst = filepath.parent / f"{filename}.mp3" + if progress_callback: + progress_callback(processed, total, filepath.name, + "local_copy", current_kbps, 0) + # Rimuovi originale solo se ha estensione diversa (altrimenti + # verrà sovrascritto dalla copia) + try: + if filepath.exists() and filepath.resolve() != dst.resolve(): + filepath.unlink() + except Exception: + pass + ok = _copy_or_convert_to_mp3(selected_path, dst, temp_dir) + if ok: + new_kbps = get_bitrate(dst) + _mark_done(done_file, filename) + _cleanup_temp(temp_dir) + processed += 1 + if progress_callback: + progress_callback(processed, total, filepath.name, + "local_upgraded", current_kbps, new_kbps) + continue + # Copia fallita: fallback YouTube (non marchiamo done) + _cleanup_temp(temp_dir) + + # ------------------------------------------------------------------ + # 2) YOUTUBE FALLBACK (comportamento originale) + # ------------------------------------------------------------------ if progress_callback: progress_callback(processed, total, filepath.name, "searching", current_kbps, 0) diff --git a/tests/test_upgrader_archive.py b/tests/test_upgrader_archive.py new file mode 100644 index 0000000..bb6e820 --- /dev/null +++ b/tests/test_upgrader_archive.py @@ -0,0 +1,169 @@ +"""Test unitari per le funzioni di archive-lookup in core.upgrader. + +Focalizzati su _normalize_stem, _scan_archive, _find_candidates. +Non testiamo il flow end-to-end di upgrade_folder (richiederebbe mock +di subprocess/yt-dlp/ffmpeg — out of scope, coperto da smoke test manuale). +""" + +from __future__ import annotations + +from pathlib import Path +from unittest.mock import patch + +import pytest + +from core import upgrader + + +# ============================================================ +# _normalize_stem +# ============================================================ +class TestNormalizeStem: + def test_lowercase_and_dedup(self): + # "Hot" e "hot" → un solo token; token duplicati collassano nel set + got = upgrader._normalize_stem("Hot Sauce hot sauce") + assert got == {"hot", "sauce"} + + def test_filters_short_tokens(self): + # Token di lunghezza < 3 vengono scartati (di, il, a, b, cd...) + got = upgrader._normalize_stem("A B cd Boombox") + assert "a" not in got + assert "b" not in got + assert "cd" not in got + assert "boombox" in got + + def test_removes_punctuation(self): + # Trattini, virgole, parentesi, apostrofi → tutti sostituiti con spazi + got = upgrader._normalize_stem("Artist - Title (Extended Mix)") + assert got == {"artist", "title", "extended", "mix"} + + def test_empty_input(self): + assert upgrader._normalize_stem("") == set() + assert upgrader._normalize_stem(" ") == set() + + def test_only_short_tokens_returns_empty(self): + assert upgrader._normalize_stem("a b c d") == set() + + +# ============================================================ +# _scan_archive +# ============================================================ +class TestScanArchive: + def test_recursive_scan(self, tmp_path: Path): + # Crea albero: top-level + sub/ + sub/sub2/ + (tmp_path / "Artist - Song.mp3").touch() + (tmp_path / "sub").mkdir() + (tmp_path / "sub" / "Second Track.mp3").touch() + (tmp_path / "sub" / "sub2").mkdir() + (tmp_path / "sub" / "sub2" / "Deep One.m4a").touch() + # File non-audio devono essere ignorati + (tmp_path / "readme.txt").touch() + + index = upgrader._scan_archive(str(tmp_path)) + # 3 entries totali (una per ogni file audio, ognuna con token distinti) + assert len(index) == 3 + # Verifica che le path siano riferite ai file corretti + all_paths = [p for paths in index.values() for p in paths] + names = sorted(p.name for p in all_paths) + assert names == ["Artist - Song.mp3", "Deep One.m4a", "Second Track.mp3"] + + def test_case_insensitive_extensions(self, tmp_path: Path): + (tmp_path / "one.MP3").touch() + (tmp_path / "two.Mp3").touch() + (tmp_path / "three.WAV").touch() + (tmp_path / "four.FLAC").touch() + (tmp_path / "five.txt").touch() # non-audio: ignorato + index = upgrader._scan_archive(str(tmp_path)) + all_paths = [p for paths in index.values() for p in paths] + assert len(all_paths) == 4 # 4 file audio, 1 skippato + + def test_duplicates_appended(self, tmp_path: Path): + # Due file con stesso token set → stessa key, due path + (tmp_path / "sub1").mkdir() + (tmp_path / "sub2").mkdir() + (tmp_path / "sub1" / "Hot Sauce.mp3").touch() + (tmp_path / "sub2" / "Hot Sauce.mp3").touch() + index = upgrader._scan_archive(str(tmp_path)) + assert len(index) == 1 + # Una sola chiave, con 2 path + paths = list(index.values())[0] + assert len(paths) == 2 + + def test_nonexistent_dir_returns_empty(self, tmp_path: Path): + assert upgrader._scan_archive(str(tmp_path / "nope")) == {} + + def test_empty_stems_skipped(self, tmp_path: Path): + # File il cui stem produce zero token utili (solo caratteri corti) → skip + (tmp_path / "a.mp3").touch() + (tmp_path / "Real Track Name.mp3").touch() + index = upgrader._scan_archive(str(tmp_path)) + # Solo "Real Track Name.mp3" ha token >=3 chars + all_paths = [p for paths in index.values() for p in paths] + assert len(all_paths) == 1 + assert all_paths[0].name == "Real Track Name.mp3" + + +# ============================================================ +# _find_candidates +# ============================================================ +class TestFindCandidates: + def _fake_index(self, tmp_path: Path, entries: list) -> dict: + """Helper: crea file (touch) e ritorna un index manuale.""" + index: dict = {} + for name in entries: + p = tmp_path / name + p.touch() + tokens = upgrader._normalize_stem(p.stem) + key = upgrader._tokens_to_key(tokens) + index.setdefault(key, []).append(p) + return index + + def test_ranking_bitrate_desc_then_similarity_desc(self, tmp_path: Path): + # 3 candidati: variamo bitrate + similarity per verificare l'ordine + # a: sim alta (0.75), bitrate basso (128) + # b: sim media (0.5), bitrate alto (320) + # c: sim alta (0.75), bitrate medio (192) + # Ordine atteso: b(320) > c(192, sim 0.75) > a(128) + idx = self._fake_index(tmp_path, [ + "Artist - Hot Sauce.mp3", # a: 3 token comuni su 4 = 0.75 + "Different Song Boombox.mp3", # b: 1 su 5 = 0.2 (troppo bassa, filtrato) + "Artist Hot Sauce Extended.mp3", # c: 3 su 4 = 0.75 + ]) + # Target: "Artist Hot Sauce" → tokens = {artist, hot, sauce} + target = "Artist Hot Sauce" + + # Mock bitrate: mappa nome → kbps + def fake_bitrate(p): + n = Path(p).name + if n == "Artist - Hot Sauce.mp3": return 128 + if n == "Different Song Boombox.mp3": return 320 + if n == "Artist Hot Sauce Extended.mp3": return 192 + return 0 + + with patch.object(upgrader, "get_bitrate", side_effect=fake_bitrate): + results = upgrader._find_candidates(target, idx, min_similarity=0.3) + + # "b" viene filtrato (sim 1/5 = 0.2 < 0.3 min); restano a e c + assert len(results) == 2 + # c (bitrate 192) prima di a (bitrate 128) + assert results[0][0].name == "Artist Hot Sauce Extended.mp3" + assert results[1][0].name == "Artist - Hot Sauce.mp3" + + def test_below_threshold_filtered(self, tmp_path: Path): + # Un solo candidato con similarity ~ 1/5 = 0.2 → sotto default 0.5 → escluso + idx = self._fake_index(tmp_path, [ + "Foo Bar Baz Qux Extra.mp3", + ]) + with patch.object(upgrader, "get_bitrate", return_value=320): + results = upgrader._find_candidates("Hot Sauce", idx) # solo "hot"/"sauce" + assert results == [] + + def test_exact_match_returned(self, tmp_path: Path): + idx = self._fake_index(tmp_path, [ + "Hot Sauce Extended.mp3", + ]) + with patch.object(upgrader, "get_bitrate", return_value=320): + results = upgrader._find_candidates("Hot Sauce Extended", idx) + assert len(results) == 1 + assert results[0][1] == 1.0 # perfect Jaccard + assert results[0][2] == 320