sourav-das/stem-separator
3
1import json2import logging3import os4import re5import shutil6import socket7import ssl8import tempfile9import time10from dataclasses import asdict, dataclass, field11from pathlib import Path12from urllib.parse import parse_qs, quote, urlparse13from urllib.request import Request, urlopen14 15from yt_dlp import YoutubeDL16 17from backend import file_manager18 19logger = logging.getLogger(__name__)20 21USER_AGENT = (22 "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "23 "AppleWebKit/537.36 (KHTML, like Gecko) "24 "Chrome/122.0.0.0 Safari/537.36"25)26 27IMPORT_RETRY_COUNT = int(os.getenv("IMPORT_RETRY_COUNT", "3"))28IMPORT_RETRY_BASE_DELAY_MS = int(os.getenv("IMPORT_RETRY_BASE_DELAY_MS", "800"))29IMPORT_FORCE_IPV4 = os.getenv("IMPORT_FORCE_IPV4", "1").lower() not in {"0", "false", "no"}30IMPORT_DIAGNOSTICS_ENABLED = os.getenv("IMPORT_DIAGNOSTICS_ENABLED", "1").lower() not in {"0", "false", "no"}31 32DEFAULT_DIAG_HOSTS = [33 "www.youtube.com",34 "youtube.com",35 "music.youtube.com",36 "youtu.be",37 "google.com",38]39 40 41@dataclass42class ImportedTrack:43 job_id: str44 filename: str45 source_url: str46 resolved_url: str | None47 title: str48 platform: str49 50 51@dataclass52class ImportFailureContext:53 stage: str54 host: str | None = None55 retryable: bool = False56 diagnostics: dict | None = None57 attempts: list[dict] = field(default_factory=list)58 59 60class SourceImportError(Exception):61 def __init__(62 self,63 message: str,64 *,65 stage: str = "download",66 host: str | None = None,67 retryable: bool = False,68 diagnostics: dict | None = None,69 attempts: list[dict] | None = None,70 ):71 super().__init__(message)72 self.context = ImportFailureContext(73 stage=stage,74 host=host,75 retryable=retryable,76 diagnostics=diagnostics,77 attempts=attempts or [],78 )79 80 81def import_source(url: str) -> ImportedTrack:82 normalized_url = normalize_url(url)83 platform = classify_platform(normalized_url)84 job_id = file_manager.create_job()85 job_dir = file_manager.get_job_dir(job_id)86 87 try:88 if platform in {"youtube", "ytmusic"}:89 title, display_filename, import_diag = download_youtube_source(90 normalized_url,91 job_dir,92 )93 metadata = {94 "source_kind": platform,95 "source_url": normalized_url,96 "resolved_url": normalized_url,97 "title": title,98 "filename": display_filename,99 "import_diagnostics": import_diag,100 }101 file_manager.save_job_metadata(job_id, metadata)102 return ImportedTrack(103 job_id=job_id,104 filename=display_filename,105 source_url=normalized_url,106 resolved_url=normalized_url,107 title=title,108 platform=platform,109 )110 111 track_title, primary_artist = fetch_spotify_track_metadata(normalized_url)112 search_query = f"{track_title} {primary_artist} audio"113 resolved_url, search_diag = resolve_youtube_search(search_query)114 title, display_filename, download_diag = download_youtube_source(115 resolved_url,116 job_dir,117 title_hint=f"{track_title} - {primary_artist}",118 )119 metadata = {120 "source_kind": "spotify",121 "source_url": normalized_url,122 "resolved_url": resolved_url,123 "title": title,124 "filename": display_filename,125 "spotify_track_title": track_title,126 "spotify_primary_artist": primary_artist,127 "import_diagnostics": {128 "search": search_diag,129 "download": download_diag,130 },131 }132 file_manager.save_job_metadata(job_id, metadata)133 return ImportedTrack(134 job_id=job_id,135 filename=display_filename,136 source_url=normalized_url,137 resolved_url=resolved_url,138 title=title,139 platform="spotify",140 )141 except Exception:142 file_manager.delete_job(job_id)143 raise144 145 146def normalize_url(url: str) -> str:147 value = url.strip()148 if not value:149 raise SourceImportError(150 "Paste a YouTube, YouTube Music, or Spotify track link",151 stage="validation",152 )153 return value154 155 156def classify_platform(url: str) -> str:157 parsed = urlparse(url)158 host = parsed.netloc.lower().removeprefix("www.")159 path = parsed.path160 query = parse_qs(parsed.query)161 162 if host in {"youtube.com", "m.youtube.com", "youtu.be"}:163 if "list" in query:164 raise SourceImportError("Playlist links are not supported yet", stage="validation")165 if host == "youtu.be":166 return "youtube"167 if path.startswith("/watch") or path.startswith("/shorts/"):168 return "youtube"169 if path.startswith("/playlist") or path.startswith("/channel/") or path.startswith("/@"):170 raise SourceImportError("Only single YouTube video links are supported", stage="validation")171 raise SourceImportError("Unsupported YouTube link", stage="validation")172 173 if host == "music.youtube.com":174 if "list" in query:175 raise SourceImportError("Playlist links are not supported yet", stage="validation")176 if path.startswith("/watch"):177 return "ytmusic"178 raise SourceImportError("Only single YouTube Music track links are supported", stage="validation")179 180 if host == "open.spotify.com":181 parts = [part for part in path.split("/") if part]182 if len(parts) >= 2 and parts[0] == "track":183 return "spotify"184 if parts and parts[0] in {"album", "playlist", "artist", "show", "episode"}:185 raise SourceImportError("Only single Spotify track links are supported", stage="validation")186 raise SourceImportError("Unsupported Spotify link", stage="validation")187 188 raise SourceImportError(189 "Unsupported link. Use YouTube, YouTube Music, or Spotify track URLs",190 stage="validation",191 )192 193 194def download_youtube_source(195 source_url: str,196 job_dir: Path,197 title_hint: str | None = None,198) -> tuple[str, str, dict]:199 host = extract_host(source_url)200 attempts: list[dict] = []201 preflight = run_host_diagnostic(host) if host else None202 203 retryable_error: SourceImportError | None = None204 205 for attempt_index in range(1, IMPORT_RETRY_COUNT + 1):206 temp_dir = Path(tempfile.mkdtemp(prefix="import-", dir=str(job_dir)))207 attempt_started = time.perf_counter()208 209 try:210 options = build_ytdlp_options(temp_dir)211 with YoutubeDL(options) as ydl:212 info = ydl.extract_info(source_url, download=True)213 214 output_path = find_downloaded_audio(temp_dir)215 content = output_path.read_bytes()216 display_title = sanitize_title(title_hint or info.get("title") or output_path.stem)217 display_filename = f"{display_title}.wav"218 file_manager.save_imported_audio(job_dir.name, display_filename, content)219 220 attempts.append(221 {222 "attempt": attempt_index,223 "status": "success",224 "duration_seconds": time.perf_counter() - attempt_started,225 }226 )227 return display_title, display_filename, {228 "stage": "download",229 "preflight": preflight,230 "attempts": attempts,231 }232 except SourceImportError as exc:233 attempts.append(build_attempt_record(attempt_index, attempt_started, exc))234 retryable_error = exc if exc.context.retryable else None235 if not exc.context.retryable or attempt_index == IMPORT_RETRY_COUNT:236 attach_and_log_diagnostics(exc, source_url, host, preflight, attempts)237 raise exc238 backoff_sleep(attempt_index)239 except Exception as exc:240 wrapped = wrap_external_error(exc, stage="download", host=host)241 attempts.append(build_attempt_record(attempt_index, attempt_started, wrapped))242 retryable_error = wrapped if wrapped.context.retryable else None243 if not wrapped.context.retryable or attempt_index == IMPORT_RETRY_COUNT:244 attach_and_log_diagnostics(wrapped, source_url, host, preflight, attempts)245 raise wrapped246 backoff_sleep(attempt_index)247 finally:248 shutil.rmtree(temp_dir, ignore_errors=True)249 250 if retryable_error is not None:251 attach_and_log_diagnostics(retryable_error, source_url, host, preflight, attempts)252 raise retryable_error253 254 raise SourceImportError(255 "Temporary YouTube connectivity issue, please retry",256 stage="download",257 host=host,258 retryable=True,259 diagnostics=preflight,260 attempts=attempts,261 )262 263 264def resolve_youtube_search(query: str) -> tuple[str, dict]:265 search_url = "https://www.youtube.com/results"266 host = extract_host(search_url)267 preflight = run_host_diagnostic(host) if host else None268 attempts: list[dict] = []269 270 for attempt_index in range(1, IMPORT_RETRY_COUNT + 1):271 attempt_started = time.perf_counter()272 try:273 options = {274 "quiet": True,275 "no_warnings": True,276 "extract_flat": "in_playlist",277 "noplaylist": True,278 }279 if IMPORT_FORCE_IPV4:280 options["source_address"] = "0.0.0.0"281 282 with YoutubeDL(options) as ydl:283 info = ydl.extract_info(f"ytsearch1:{query}", download=False)284 except Exception as exc:285 wrapped = wrap_external_error(exc, stage="search", host=host)286 attempts.append(build_attempt_record(attempt_index, attempt_started, wrapped))287 if not wrapped.context.retryable or attempt_index == IMPORT_RETRY_COUNT:288 attach_and_log_diagnostics(wrapped, query, host, preflight, attempts)289 raise wrapped290 backoff_sleep(attempt_index)291 continue292 293 entries = info.get("entries") or []294 if not entries:295 error = SourceImportError(296 "No matching YouTube source was found for this Spotify track",297 stage="search",298 host=host,299 retryable=False,300 )301 attempts.append(build_attempt_record(attempt_index, attempt_started, error))302 attach_and_log_diagnostics(error, query, host, preflight, attempts)303 raise error304 305 entry = entries[0]306 resolved_url = entry.get("webpage_url") or entry.get("url")307 if resolved_url and not str(resolved_url).startswith("http"):308 video_id = entry.get("id") or resolved_url309 resolved_url = f"https://www.youtube.com/watch?v={video_id}"310 if not resolved_url:311 error = SourceImportError(312 "Resolved YouTube match did not include a downloadable URL",313 stage="search",314 host=host,315 retryable=False,316 )317 attempts.append(build_attempt_record(attempt_index, attempt_started, error))318 attach_and_log_diagnostics(error, query, host, preflight, attempts)319 raise error320 321 attempts.append(322 {323 "attempt": attempt_index,324 "status": "success",325 "duration_seconds": time.perf_counter() - attempt_started,326 }327 )328 return resolved_url, {"stage": "search", "preflight": preflight, "attempts": attempts}329 330 raise SourceImportError(331 "Temporary YouTube connectivity issue, please retry",332 stage="search",333 host=host,334 retryable=True,335 diagnostics=preflight,336 attempts=attempts,337 )338 339 340def fetch_spotify_track_metadata(url: str) -> tuple[str, str]:341 oembed_url = f"https://open.spotify.com/oembed?url={quote(url, safe='')}"342 request = Request(oembed_url, headers={"User-Agent": USER_AGENT})343 344 try:345 with urlopen(request, timeout=20) as response:346 payload = json.loads(response.read().decode("utf-8"))347 title = payload.get("title", "").strip()348 artist = payload.get("author_name", "").strip()349 if title and artist:350 return title, artist351 except Exception:352 pass353 354 html = fetch_text(url)355 title = first_match(356 html,357 [358 r'<meta property="og:title" content="([^"]+)"',359 r"<title>([^<]+)</title>",360 ],361 )362 artist = first_match(363 html,364 [365 r'<meta name="music:musician_description" content="([^"]+)"',366 r'"artists"\s*:\s*\[\s*\{\s*"name"\s*:\s*"([^"]+)"',367 r'"byArtist"\s*:\s*\{\s*"name"\s*:\s*"([^"]+)"',368 ],369 )370 371 if not title or not artist:372 raise SourceImportError(373 "Could not read Spotify track metadata from the public page",374 stage="metadata",375 )376 377 return clean_spotify_title(title), artist.strip()378 379 380def fetch_text(url: str) -> str:381 request = Request(url, headers={"User-Agent": USER_AGENT})382 with urlopen(request, timeout=20) as response:383 return response.read().decode("utf-8", errors="ignore")384 385 386def find_downloaded_audio(temp_dir: Path) -> Path:387 audio_files = sorted(388 [389 path390 for path in temp_dir.iterdir()391 if path.is_file()392 and path.suffix.lower() in {".wav", ".mp3", ".m4a", ".aac", ".flac", ".opus", ".ogg"}393 ],394 key=lambda path: path.stat().st_mtime,395 reverse=True,396 )397 if not audio_files:398 raise SourceImportError(399 "Downloaded source did not produce a playable audio file",400 stage="postprocess",401 )402 return audio_files[0]403 404 405def build_ytdlp_options(temp_dir: Path) -> dict:406 options = {407 "format": "bestaudio/best",408 "paths": {"home": str(temp_dir)},409 "outtmpl": {"default": "downloaded.%(ext)s"},410 "quiet": True,411 "no_warnings": True,412 "noplaylist": True,413 "extract_flat": False,414 "retries": 1,415 "fragment_retries": 1,416 "postprocessors": [417 {418 "key": "FFmpegExtractAudio",419 "preferredcodec": "wav",420 }421 ],422 }423 if IMPORT_FORCE_IPV4:424 options["source_address"] = "0.0.0.0"425 return options426 427 428def run_network_diagnostics() -> list[dict]:429 return [run_host_diagnostic(host) for host in DEFAULT_DIAG_HOSTS]430 431 432def run_host_diagnostic(host: str, port: int = 443) -> dict:433 result: dict = {"host": host, "port": port}434 435 dns_started = time.perf_counter()436 try:437 info = socket.getaddrinfo(438 host,439 port,440 family=socket.AF_INET if IMPORT_FORCE_IPV4 else socket.AF_UNSPEC,441 type=socket.SOCK_STREAM,442 )443 ip = info[0][4][0]444 result["dns_ok"] = True445 result["resolved_ip"] = ip446 except Exception as exc:447 result["dns_ok"] = False448 result["dns_error"] = str(exc)449 result["dns_seconds"] = time.perf_counter() - dns_started450 return result451 result["dns_seconds"] = time.perf_counter() - dns_started452 453 connect_started = time.perf_counter()454 try:455 family, socktype, proto, _, sockaddr = info[0]456 with socket.socket(family, socktype, proto) as raw_sock:457 raw_sock.settimeout(5)458 raw_sock.connect(sockaddr)459 context = ssl.create_default_context()460 with context.wrap_socket(raw_sock, server_hostname=host):461 pass462 result["https_ok"] = True463 except Exception as exc:464 result["https_ok"] = False465 result["https_error"] = str(exc)466 result["https_seconds"] = time.perf_counter() - connect_started467 return result468 469 470def wrap_external_error(exc: Exception, *, stage: str, host: str | None) -> SourceImportError:471 message = str(exc)472 lower = message.lower()473 474 if "failed to resolve" in lower or "temporary failure in name resolution" in lower or "no address associated with hostname" in lower:475 return SourceImportError(476 "Could not resolve YouTube host from the Space runtime",477 stage="dns",478 host=host,479 retryable=True,480 )481 482 if any(token in lower for token in ["timed out", "timeout", "connection reset", "network is unreachable", "transporterror", "ssl"]):483 return SourceImportError(484 "Temporary YouTube connectivity issue, please retry",485 stage=stage,486 host=host,487 retryable=True,488 )489 490 return SourceImportError(491 f"YouTube {stage} request failed after retries",492 stage=stage,493 host=host,494 retryable=False,495 )496 497 498def attach_and_log_diagnostics(499 error: SourceImportError,500 target: str,501 host: str | None,502 preflight: dict | None,503 attempts: list[dict],504) -> None:505 error.context.host = host or error.context.host506 error.context.diagnostics = preflight507 error.context.attempts = attempts508 if IMPORT_DIAGNOSTICS_ENABLED:509 logger.warning(510 "source_import_failure %s",511 json.dumps(512 {513 "target": target,514 "message": str(error),515 "context": asdict(error.context),516 }517 ),518 )519 520 521def build_attempt_record(522 attempt_index: int,523 attempt_started: float,524 error: SourceImportError,525) -> dict:526 return {527 "attempt": attempt_index,528 "status": "error",529 "duration_seconds": time.perf_counter() - attempt_started,530 "stage": error.context.stage,531 "retryable": error.context.retryable,532 "message": str(error),533 }534 535 536def backoff_sleep(attempt_index: int) -> None:537 delay = (IMPORT_RETRY_BASE_DELAY_MS / 1000.0) * (2 ** (attempt_index - 1))538 time.sleep(delay)539 540 541def extract_host(url: str) -> str | None:542 parsed = urlparse(url)543 return parsed.netloc or None544 545 546def first_match(text: str, patterns: list[str]) -> str | None:547 for pattern in patterns:548 match = re.search(pattern, text, re.IGNORECASE)549 if match:550 return unescape_html(match.group(1))551 return None552 553 554def sanitize_title(value: str) -> str:555 clean = re.sub(r"[\\/:*?\"<>|]+", " ", value)556 clean = re.sub(r"\s+", " ", clean).strip().strip(".")557 return clean[:120] or "Imported Track"558 559 560def clean_spotify_title(value: str) -> str:561 title = value.replace(" | Spotify", "").strip()562 if " - song and lyrics by " in title.lower():563 title = re.split(r"\s+-\s+song and lyrics by\s+", title, flags=re.IGNORECASE)[0]564 return title.strip()565 566 567def unescape_html(value: str) -> str:568 return (569 value.replace("&", "&")570 .replace(""", '"')571 .replace("'", "'")572 .replace("'", "'")573 )574