Add YtdlpClient: shared JSON invocation layer for yt-dlp
Every metadata call site previously built its own command list. The new ytsage/core/ytsage_client.py centralizes binary resolution (managed, verified binary only), cookie/proxy args from ConfigManager (with session-state overrides), utf-8 output handling, process-group-safe timeouts, and a short-TTL cache for flat/search results. Provides fetch_video_info, fetch_flat_info, fetch_flat_entries (with -I range pagination), search (ytsearchN:), and fetch_account_feed (youtube.com/feed/subscriptions with cookies) plus a generic YtdlpWorker QThread. AnalysisThread now delegates its subprocess execution to the shared runner; its signal surface is unchanged. Groundwork for the SageTube watch/search/browse/feed features. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,262 @@
|
||||
"""
|
||||
YtdlpClient - shared yt-dlp invocation layer
|
||||
============================================
|
||||
|
||||
Single place that knows how to build and run yt-dlp commands for metadata
|
||||
purposes (analysis, search, channel/playlist browsing, subscription feeds).
|
||||
Every caller previously built its own command list; new features should go
|
||||
through this client so cookies, proxies, timeouts and process-group cleanup
|
||||
behave identically everywhere.
|
||||
|
||||
Downloads keep using DownloadThread (streaming progress parsing); this client
|
||||
is for JSON-returning invocations.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import signal
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, Dict, List, Optional, Tuple
|
||||
|
||||
from PySide6.QtCore import QThread, Signal
|
||||
|
||||
from .ytsage_yt_dlp import get_yt_dlp_path
|
||||
from ..utils.ytsage_config_manager import ConfigManager
|
||||
from ..utils.ytsage_constants import SUBPROCESS_CREATIONFLAGS
|
||||
from ..utils.ytsage_logger import logger
|
||||
|
||||
# Flat-entry results are cached briefly so tab switches and pagination don't
|
||||
# hammer YouTube with identical requests
|
||||
_CACHE_TTL_SECONDS = 600
|
||||
_cache: Dict[Tuple[str, ...], Tuple[float, Any]] = {}
|
||||
_cache_lock = threading.Lock()
|
||||
|
||||
|
||||
class YtdlpNotInstalledError(FileNotFoundError):
|
||||
"""Raised when neither the managed binary nor an opted-in system yt-dlp exists."""
|
||||
|
||||
|
||||
def run_ytdlp_capture(cmd: List[str], timeout: int) -> subprocess.CompletedProcess:
|
||||
"""Run yt-dlp capturing output, killing the whole process group on timeout.
|
||||
|
||||
subprocess.run() only kills the direct child on TimeoutExpired, leaking
|
||||
grandchildren (deno); it also loses any stderr produced before the
|
||||
timeout. Use Popen with a new session so the entire group can be reaped,
|
||||
and surface the partial stderr on the raised exception.
|
||||
"""
|
||||
popen_kwargs: Dict[str, Any] = {
|
||||
"stdout": subprocess.PIPE,
|
||||
"stderr": subprocess.PIPE,
|
||||
"text": True,
|
||||
"encoding": "utf-8",
|
||||
"errors": "replace",
|
||||
}
|
||||
if sys.platform == "win32":
|
||||
popen_kwargs["creationflags"] = SUBPROCESS_CREATIONFLAGS
|
||||
else:
|
||||
popen_kwargs["start_new_session"] = True
|
||||
|
||||
proc = subprocess.Popen(cmd, **popen_kwargs)
|
||||
try:
|
||||
stdout, stderr = proc.communicate(timeout=timeout)
|
||||
except subprocess.TimeoutExpired as exc:
|
||||
if sys.platform == "win32":
|
||||
subprocess.run(
|
||||
["taskkill", "/F", "/T", "/PID", str(proc.pid)],
|
||||
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
|
||||
creationflags=SUBPROCESS_CREATIONFLAGS,
|
||||
)
|
||||
else:
|
||||
try:
|
||||
os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
|
||||
except (ProcessLookupError, PermissionError, OSError):
|
||||
pass
|
||||
stdout, stderr = proc.communicate()
|
||||
exc.stdout, exc.stderr = stdout, stderr
|
||||
raise
|
||||
return subprocess.CompletedProcess(cmd, proc.returncode, stdout, stderr)
|
||||
|
||||
|
||||
class YtdlpClient:
|
||||
"""Builds and runs JSON-returning yt-dlp commands with the app's
|
||||
cookie/proxy configuration applied."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
cookie_file_path: Optional[str] = None,
|
||||
browser_cookies: Optional[str] = None,
|
||||
use_config_auth: bool = True,
|
||||
) -> None:
|
||||
"""
|
||||
Args:
|
||||
cookie_file_path / browser_cookies: explicit overrides for session
|
||||
state that differs from the persisted config.
|
||||
use_config_auth: when True (default) and no override is given,
|
||||
cookie and proxy options are read from ConfigManager.
|
||||
"""
|
||||
self._cookie_file_override = cookie_file_path
|
||||
self._browser_cookies_override = browser_cookies
|
||||
self._use_config_auth = use_config_auth
|
||||
|
||||
# ------------------------------------------------------------- commands
|
||||
|
||||
def resolve_binary(self) -> str:
|
||||
path = get_yt_dlp_path()
|
||||
if str(path) == "yt-dlp":
|
||||
raise YtdlpNotInstalledError("yt-dlp is not installed - run the yt-dlp setup first")
|
||||
return str(path)
|
||||
|
||||
def _auth_args(self) -> List[str]:
|
||||
args: List[str] = []
|
||||
cookie_file = self._cookie_file_override
|
||||
browser_cookies = self._browser_cookies_override
|
||||
|
||||
if cookie_file is None and browser_cookies is None and self._use_config_auth:
|
||||
if ConfigManager.get("cookie_active"):
|
||||
if ConfigManager.get("cookie_source") == "file":
|
||||
saved = ConfigManager.get("cookie_file_path")
|
||||
if saved and Path(saved).exists():
|
||||
cookie_file = str(saved)
|
||||
else:
|
||||
browser = ConfigManager.get("cookie_browser")
|
||||
profile = ConfigManager.get("cookie_browser_profile")
|
||||
if browser:
|
||||
browser_cookies = f"{browser}:{profile}" if profile else browser
|
||||
|
||||
if cookie_file:
|
||||
args.extend(["--cookies", str(cookie_file)])
|
||||
elif browser_cookies:
|
||||
args.extend(["--cookies-from-browser", browser_cookies])
|
||||
|
||||
if self._use_config_auth:
|
||||
proxy = ConfigManager.get("proxy_url")
|
||||
geo_proxy = ConfigManager.get("geo_proxy_url")
|
||||
if proxy:
|
||||
args.extend(["--proxy", str(proxy)])
|
||||
if geo_proxy:
|
||||
args.extend(["--geo-verification-proxy", str(geo_proxy)])
|
||||
return args
|
||||
|
||||
def build_base_cmd(self) -> List[str]:
|
||||
return [self.resolve_binary(), "--no-warnings", "--no-color"] + self._auth_args()
|
||||
|
||||
def run_json(self, extra_args: List[str], timeout: int = 60) -> Any:
|
||||
"""Run yt-dlp with --dump-single-json semantics and parse stdout."""
|
||||
cmd = self.build_base_cmd() + extra_args
|
||||
logger.debug(f"YtdlpClient executing: {cmd}")
|
||||
result = run_ytdlp_capture(cmd, timeout=timeout)
|
||||
if result.returncode != 0:
|
||||
stderr_tail = (result.stderr or "").strip()[-1000:]
|
||||
raise RuntimeError(stderr_tail or f"yt-dlp exited with code {result.returncode}")
|
||||
return json.loads(result.stdout)
|
||||
|
||||
# ------------------------------------------------------------- queries
|
||||
|
||||
def fetch_video_info(self, url: str, timeout: int = 300) -> Dict[str, Any]:
|
||||
"""Full info dict for a single video, including formats[].url."""
|
||||
return self.run_json(["--dump-single-json", url], timeout=timeout)
|
||||
|
||||
def fetch_flat_info(self, url: str, timeout: int = 300) -> Dict[str, Any]:
|
||||
"""Flat info dict (playlist/channel entries without per-video formats)."""
|
||||
return self.run_json(["--dump-single-json", "--flat-playlist", url], timeout=timeout)
|
||||
|
||||
def fetch_flat_entries(
|
||||
self, url: str, start: int = 1, end: int = 30, timeout: int = 120, use_cache: bool = True
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Flat entries start..end (1-based, inclusive) of a playlist/channel URL."""
|
||||
cache_key = ("flat", url, str(start), str(end), *self._auth_args())
|
||||
if use_cache:
|
||||
cached = _cache_get(cache_key)
|
||||
if cached is not None:
|
||||
return cached
|
||||
info = self.run_json(
|
||||
["--dump-single-json", "--flat-playlist", "-I", f"{start}:{end}", url],
|
||||
timeout=timeout,
|
||||
)
|
||||
entries = [e for e in info.get("entries", []) if e] if isinstance(info, dict) else []
|
||||
_cache_put(cache_key, entries)
|
||||
return entries
|
||||
|
||||
def search(
|
||||
self, query: str, n: int = 25, offset: int = 0, timeout: int = 120, use_cache: bool = True
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""YouTube search returning flat entries n results at a time."""
|
||||
cache_key = ("search", query, str(n), str(offset), *self._auth_args())
|
||||
if use_cache:
|
||||
cached = _cache_get(cache_key)
|
||||
if cached is not None:
|
||||
return cached
|
||||
total = offset + n
|
||||
info = self.run_json(
|
||||
[
|
||||
"--dump-single-json",
|
||||
"--flat-playlist",
|
||||
"-I",
|
||||
f"{offset + 1}:{total}",
|
||||
f"ytsearch{total}:{query}",
|
||||
],
|
||||
timeout=timeout,
|
||||
)
|
||||
entries = [e for e in info.get("entries", []) if e] if isinstance(info, dict) else []
|
||||
_cache_put(cache_key, entries)
|
||||
return entries
|
||||
|
||||
def fetch_account_feed(self, n: int = 50, timeout: int = 180) -> List[Dict[str, Any]]:
|
||||
"""The logged-in account's subscription feed. Requires active cookies."""
|
||||
return self.fetch_flat_entries(
|
||||
"https://www.youtube.com/feed/subscriptions", start=1, end=n,
|
||||
timeout=timeout, use_cache=False,
|
||||
)
|
||||
|
||||
|
||||
class YtdlpWorker(QThread):
|
||||
"""Generic worker running one YtdlpClient call off the GUI thread.
|
||||
|
||||
Usage:
|
||||
worker = YtdlpWorker(lambda c: c.search("query", 25))
|
||||
worker.result.connect(...); worker.error.connect(...)
|
||||
worker.start()
|
||||
"""
|
||||
|
||||
result = Signal(object)
|
||||
error = Signal(str)
|
||||
|
||||
def __init__(self, fn: Callable[[YtdlpClient], Any], client: Optional[YtdlpClient] = None, parent=None) -> None:
|
||||
super().__init__(parent)
|
||||
self._fn = fn
|
||||
self._client = client or YtdlpClient()
|
||||
|
||||
def run(self) -> None:
|
||||
try:
|
||||
self.result.emit(self._fn(self._client))
|
||||
except Exception as e:
|
||||
logger.exception(f"YtdlpWorker failed: {e}")
|
||||
self.error.emit(str(e))
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ cache
|
||||
|
||||
def _cache_get(key: Tuple[str, ...]) -> Optional[Any]:
|
||||
with _cache_lock:
|
||||
hit = _cache.get(key)
|
||||
if hit and time.time() - hit[0] < _CACHE_TTL_SECONDS:
|
||||
return hit[1]
|
||||
if hit:
|
||||
del _cache[key]
|
||||
return None
|
||||
|
||||
|
||||
def _cache_put(key: Tuple[str, ...], value: Any) -> None:
|
||||
with _cache_lock:
|
||||
if len(_cache) > 256:
|
||||
_cache.clear()
|
||||
_cache[key] = (time.time(), value)
|
||||
|
||||
|
||||
def clear_cache() -> None:
|
||||
with _cache_lock:
|
||||
_cache.clear()
|
||||
Reference in New Issue
Block a user