Apply targeted fixes to prodl-core

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(".")

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

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"