diff --git a/packages/prodl-core/prodl_core/downloader.py b/packages/prodl-core/prodl_core/downloader.py index 49759db..538bc23 100644 --- a/packages/prodl-core/prodl_core/downloader.py +++ b/packages/prodl-core/prodl_core/downloader.py @@ -1,5 +1,4 @@ import asyncio -import glob import os from typing import Any import yt_dlp @@ -125,14 +124,13 @@ def _extract() -> StreamInfo: if info is None: raise ValueError(f"Could not extract stream info for URL: {opts.url}") - direct_url = info.get("url") - if not direct_url: - req_dl = info.get("requested_downloads") - if req_dl and isinstance(req_dl, list) and len(req_dl) > 0: - direct_url = req_dl[0].get("url") - - if not direct_url: - raise ValueError(f"Direct stream URL not found for: {opts.url}") + # Try requested_downloads first (most reliable) + requested = info.get('requested_downloads') + if requested and len(requested) > 0: + direct_url = requested[0]['url'] + else: + # Fallback for single-format responses + direct_url = info['url'] filesize = info.get("filesize") or info.get("filesize_approx") @@ -190,29 +188,20 @@ async def download(opts: DownloadOptions) -> DownloadResult: elif opts.sb_action == "mark": ydl_opts['sponsorblock_chapter_title'] = opts.sb_categories + if opts.progress_hook: + ydl_opts['progress_hooks'] = [opts.progress_hook] + + if opts.max_filesize: + ydl_opts['max_filesize'] = opts.max_filesize + def _download() -> DownloadResult: with yt_dlp.YoutubeDL(ydl_opts) as ydl: info = ydl.extract_info(opts.url, download=True) if info is None: raise ValueError(f"Download failed for URL: {opts.url}") - video_id = info.get("id") title = str(info.get("title") or "") - - # Find output file by glob pattern - matching_files = glob.glob(f"{opts.output_dir}/{video_id}.*") - - # Exclude temporary download files or info json if any - matching_files = [ - f for f in matching_files - if not f.endswith(".part") and not f.endswith(".info.json") and not f.endswith(".ytdl") - ] - - if matching_files: - filepath = matching_files[0] - else: - # Fallback to requested download filename or expected filepath - filepath = f"{opts.output_dir}/{video_id}.mp4" + filepath = info['requested_downloads'][0]['filepath'] filename = os.path.basename(filepath) ext = os.path.splitext(filepath)[1].lstrip(".") diff --git a/packages/prodl-core/prodl_core/models.py b/packages/prodl-core/prodl_core/models.py index 9d5b2fd..6aae2e3 100644 --- a/packages/prodl-core/prodl_core/models.py +++ b/packages/prodl-core/prodl_core/models.py @@ -1,4 +1,4 @@ -from typing import Literal +from typing import Callable, Literal from pydantic import BaseModel @@ -31,6 +31,16 @@ class VideoInfo(BaseModel): class DownloadOptions(BaseModel): + """ + Options for downloading media. + + progress_hook signature called by yt-dlp: + def hook(d: dict): + d['status'] # 'downloading' | 'finished' | 'error' + d['_percent_str'] # '45.2%' + d['_speed_str'] # '1.23MiB/s' + d['eta'] # seconds remaining (int) + """ url: str format_id: str = "bestvideo[ext=mp4]+bestaudio[ext=m4a]/best" audio_only: bool = False @@ -53,6 +63,12 @@ class DownloadOptions(BaseModel): # Output output_dir: str = "/tmp/prodl_downloads" + progress_hook: Callable | None = None + max_filesize: str | None = None + + class Config: + arbitrary_types_allowed = True + class DownloadResult(BaseModel): filepath: str diff --git a/packages/prodl-core/tests/test_downloader.py b/packages/prodl-core/tests/test_downloader.py index 202b91c..8ec00fd 100644 --- a/packages/prodl-core/tests/test_downloader.py +++ b/packages/prodl-core/tests/test_downloader.py @@ -45,7 +45,27 @@ async def test_get_info_mock(): @pytest.mark.asyncio -async def test_get_stream_info_mock(): +async def test_get_stream_info_requested_downloads(): + mock_info = { + "requested_downloads": [{"url": "https://cdn.example.com/req.mp4"}], + "url": "https://cdn.example.com/fallback.mp4", + "title": "Stream Title", + "ext": "mp4", + "filesize": 12345, + } + + with patch("yt_dlp.YoutubeDL") as MockYDL: + instance = MockYDL.return_value.__enter__.return_value + instance.extract_info.return_value = mock_info + + opts = DownloadOptions(url="https://example.com/video") + stream_info = await ProDLDownloader.get_stream_info(opts) + assert stream_info.url == "https://cdn.example.com/req.mp4" + assert stream_info.title == "Stream Title" + + +@pytest.mark.asyncio +async def test_get_stream_info_fallback(): mock_info = { "url": "https://cdn.example.com/direct.mp4", "title": "Stream Title", @@ -76,8 +96,14 @@ async def test_download_mock(tmp_path): "title": "Downloaded Title", "ext": "mp4", "filesize": 13, + "requested_downloads": [ + {"filepath": filepath} + ] } + def my_hook(d): + pass + with patch("yt_dlp.YoutubeDL") as MockYDL: instance = MockYDL.return_value.__enter__.return_value instance.extract_info.return_value = mock_info @@ -89,9 +115,16 @@ async def test_download_mock(tmp_path): thumbnail=True, sponsorblock=True, sb_action="cut", + progress_hook=my_hook, + max_filesize="50M", ) res = await ProDLDownloader.download(opts) assert res.title == "Downloaded Title" assert res.filename == f"{mock_id}.mp4" assert res.filepath == filepath assert res.filesize == 13 + + # Verify ydl_opts received options + called_opts = MockYDL.call_args[0][0] + assert called_opts.get("progress_hooks") == [my_hook] + assert called_opts.get("max_filesize") == "50M"