Refactor: try to keep extra information available in all stages of the request process
This commit is contained in:
parent
d7fc66a63a
commit
459dee77db
5 changed files with 339 additions and 5 deletions
|
|
@ -35,6 +35,7 @@ from .Utils import (
|
|||
dt_delta,
|
||||
extract_info,
|
||||
extract_ytdlp_logs,
|
||||
get_extras,
|
||||
merge_dict,
|
||||
str_to_dt,
|
||||
ytdlp_reject,
|
||||
|
|
@ -365,8 +366,9 @@ class DownloadQueue(metaclass=Singleton):
|
|||
if property in entry:
|
||||
extras[f"playlist_{property}"] = entry.get(property)
|
||||
|
||||
if "thumbnail" not in etr and "youtube:" in entry.get("extractor", ""):
|
||||
extras["thumbnail"] = f"https://img.youtube.com/vi/{etr['id']}/maxresdefault.jpg"
|
||||
extractor_key = entry.get("ie_key") or entry.get("extractor_key") or entry.get("extractor") or ""
|
||||
if "thumbnail" not in etr and "youtube" in str(extractor_key).lower():
|
||||
extras["thumbnail"] = "https://img.youtube.com/vi/{id}/maxresdefault.jpg".format(**etr)
|
||||
|
||||
newItem: Item = item.new_with(url=etr.get("url") or etr.get("webpage_url"), extras=extras)
|
||||
|
||||
|
|
@ -521,7 +523,7 @@ class DownloadQueue(metaclass=Singleton):
|
|||
options=options,
|
||||
cli=item.cli,
|
||||
auto_start=item.auto_start,
|
||||
extras=item.extras,
|
||||
extras=merge_dict(item.extras, get_extras(entry)),
|
||||
)
|
||||
|
||||
try:
|
||||
|
|
@ -838,7 +840,7 @@ class DownloadQueue(metaclass=Singleton):
|
|||
cookies=item.cookies,
|
||||
template=item.template,
|
||||
msg="URL is already downloaded.",
|
||||
extras=item.extras,
|
||||
extras=merge_dict(item.extras, get_extras(entry)),
|
||||
)
|
||||
)
|
||||
|
||||
|
|
@ -892,7 +894,7 @@ class DownloadQueue(metaclass=Singleton):
|
|||
cookies=item.cookies,
|
||||
template=item.template,
|
||||
msg=log_message,
|
||||
extras=item.extras,
|
||||
extras=merge_dict(item.extras, get_extras(entry)),
|
||||
)
|
||||
)
|
||||
await self.done.put(dlInfo)
|
||||
|
|
|
|||
|
|
@ -102,6 +102,20 @@ class Item:
|
|||
"""
|
||||
return self.extras and len(self.extras) > 0
|
||||
|
||||
def add_extras(self, key: str, value: Any) -> None:
|
||||
"""
|
||||
Add an extra data to the item.
|
||||
|
||||
Args:
|
||||
key (str): The key of the extra data.
|
||||
value (Any): The value of the extra data.
|
||||
|
||||
"""
|
||||
if not self.extras:
|
||||
self.extras = {}
|
||||
|
||||
self.extras[key] = value
|
||||
|
||||
def has_cli(self) -> bool:
|
||||
"""
|
||||
Check if the item has any command options for yt-dlp associated with it.
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ import socket
|
|||
import time
|
||||
import uuid
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from email.utils import formatdate
|
||||
from functools import lru_cache, wraps
|
||||
from http.cookiejar import MozillaCookieJar
|
||||
from pathlib import Path
|
||||
|
|
@ -1847,3 +1848,78 @@ def create_cookies_file(cookies: str, file: Path | None = None) -> Path:
|
|||
load_cookies(file)
|
||||
|
||||
return file
|
||||
|
||||
|
||||
def get_thumbnail(thumbnails: list) -> str | None:
|
||||
"""
|
||||
Extract thumbnail URL from a yt-dlp entry.
|
||||
|
||||
Args:
|
||||
thumbnails (list): The list of thumbnail dictionaries from yt-dlp entry.
|
||||
|
||||
Returns:
|
||||
str | None: The thumbnail URL if available, otherwise None.
|
||||
|
||||
"""
|
||||
if not thumbnails or not isinstance(thumbnails, list):
|
||||
return None
|
||||
|
||||
def _thumb_sort_key(thumb: dict) -> tuple:
|
||||
return (
|
||||
thumb.get("preference") if thumb.get("preference") is not None else -1,
|
||||
thumb.get("width") if thumb.get("width") is not None else -1,
|
||||
thumb.get("height") if thumb.get("height") is not None else -1,
|
||||
thumb.get("id") if thumb.get("id") is not None else "",
|
||||
thumb.get("url"),
|
||||
)
|
||||
|
||||
return max(thumbnails, key=_thumb_sort_key, default=None)
|
||||
|
||||
|
||||
def get_extras(entry: dict, kind: str = "video") -> dict:
|
||||
"""
|
||||
Extract useful information from a yt-dlp entry.
|
||||
|
||||
Args:
|
||||
entry (dict): The entry data from yt-dlp.
|
||||
kind (str): The type of the item (e.g., "video", "playlist").
|
||||
|
||||
Returns:
|
||||
dict: The extracted data.
|
||||
|
||||
"""
|
||||
extras = {}
|
||||
|
||||
if not entry or not isinstance(entry, dict):
|
||||
return extras
|
||||
|
||||
if "playlist" == kind:
|
||||
for property in ("id", "title", "uploader", "uploader_id"):
|
||||
if val := entry.get(property):
|
||||
extras[f"playlist_{property}"] = val
|
||||
|
||||
if thumbnail := get_thumbnail(entry.get("thumbnails", [])):
|
||||
extras["thumbnail"] = thumbnail.get("url")
|
||||
elif thumbnail := entry.get("thumbnail"):
|
||||
extras["thumbnail"] = thumbnail
|
||||
|
||||
for property in ("uploader", "channel"):
|
||||
if val := entry.get(property):
|
||||
extras[property] = val
|
||||
|
||||
if release_in := entry.get("release_timestamp"):
|
||||
extras["release_in"] = formatdate(release_in, usegmt=True)
|
||||
|
||||
if release_in and "is_upcoming" == entry.get("live_status"):
|
||||
extras["is_live"] = release_in
|
||||
|
||||
if duration := entry.get("duration"):
|
||||
extras["duration"] = duration
|
||||
|
||||
extras["is_premiere"] = bool(entry.get("is_premiere", False))
|
||||
|
||||
extractor_key = entry.get("ie_key") or entry.get("extractor_key") or entry.get("extractor") or ""
|
||||
if "thumbnail" not in extras and "youtube" in str(extractor_key).lower():
|
||||
extras["thumbnail"] = "https://img.youtube.com/vi/{id}/maxresdefault.jpg".format(**entry)
|
||||
|
||||
return extras
|
||||
|
|
|
|||
|
|
@ -329,3 +329,57 @@ class TestItemDTO:
|
|||
|
||||
mock_presets.return_value.get.assert_called_once_with("nonexistent")
|
||||
assert result is None
|
||||
|
||||
|
||||
class TestItemAddExtras:
|
||||
def test_add_extras_to_empty_dict(self):
|
||||
"""Test adding extras when extras dict is empty."""
|
||||
item = Item(url="https://example.com")
|
||||
item.extras = {}
|
||||
|
||||
item.add_extras("key1", "value1")
|
||||
|
||||
assert item.extras["key1"] == "value1"
|
||||
|
||||
def test_add_extras_when_extras_is_none(self):
|
||||
"""Test adding extras when extras is None."""
|
||||
item = Item(url="https://example.com")
|
||||
item.extras = None
|
||||
|
||||
item.add_extras("key1", "value1")
|
||||
|
||||
assert item.extras == {"key1": "value1"}
|
||||
|
||||
def test_add_extras_to_existing_dict(self):
|
||||
"""Test adding extras to an existing extras dict."""
|
||||
item = Item(url="https://example.com", extras={"existing": "data"})
|
||||
|
||||
item.add_extras("new_key", "new_value")
|
||||
|
||||
assert item.extras["existing"] == "data"
|
||||
assert item.extras["new_key"] == "new_value"
|
||||
|
||||
def test_add_extras_overwrites_existing_key(self):
|
||||
"""Test that adding extras overwrites existing keys."""
|
||||
item = Item(url="https://example.com", extras={"key1": "old_value"})
|
||||
|
||||
item.add_extras("key1", "new_value")
|
||||
|
||||
assert item.extras["key1"] == "new_value"
|
||||
|
||||
def test_add_extras_with_various_types(self):
|
||||
"""Test adding extras with various data types."""
|
||||
item = Item(url="https://example.com")
|
||||
item.extras = {}
|
||||
|
||||
item.add_extras("string", "value")
|
||||
item.add_extras("number", 42)
|
||||
item.add_extras("boolean", True)
|
||||
item.add_extras("list", [1, 2, 3])
|
||||
item.add_extras("dict", {"nested": "data"})
|
||||
|
||||
assert item.extras["string"] == "value"
|
||||
assert item.extras["number"] == 42
|
||||
assert item.extras["boolean"] is True
|
||||
assert item.extras["list"] == [1, 2, 3]
|
||||
assert item.extras["dict"] == {"nested": "data"}
|
||||
|
|
|
|||
|
|
@ -2487,3 +2487,191 @@ class TestMoveFile:
|
|||
|
||||
# Original file should still exist
|
||||
assert test_file.exists()
|
||||
|
||||
|
||||
class TestGetThumbnail:
|
||||
def test_returns_none_for_empty_list(self):
|
||||
"""Test that None is returned for an empty thumbnail list."""
|
||||
from app.library.Utils import get_thumbnail
|
||||
|
||||
assert get_thumbnail([]) is None
|
||||
|
||||
def test_returns_none_for_non_list(self):
|
||||
"""Test that None is returned for non-list input."""
|
||||
from app.library.Utils import get_thumbnail
|
||||
|
||||
assert get_thumbnail(None) is None
|
||||
assert get_thumbnail("not a list") is None
|
||||
assert get_thumbnail({"not": "list"}) is None
|
||||
|
||||
def test_returns_highest_preference_thumbnail(self):
|
||||
"""Test that the thumbnail with highest preference is returned."""
|
||||
from app.library.Utils import get_thumbnail
|
||||
|
||||
thumbnails = [
|
||||
{"url": "low.jpg", "preference": 1, "width": 100, "height": 100},
|
||||
{"url": "high.jpg", "preference": 10, "width": 200, "height": 200},
|
||||
{"url": "medium.jpg", "preference": 5, "width": 150, "height": 150},
|
||||
]
|
||||
|
||||
result = get_thumbnail(thumbnails)
|
||||
assert result == {"url": "high.jpg", "preference": 10, "width": 200, "height": 200}
|
||||
|
||||
def test_returns_highest_width_when_preference_equal(self):
|
||||
"""Test that the thumbnail with highest width is returned when preference is equal."""
|
||||
from app.library.Utils import get_thumbnail
|
||||
|
||||
thumbnails = [
|
||||
{"url": "small.jpg", "preference": 1, "width": 100, "height": 100},
|
||||
{"url": "large.jpg", "preference": 1, "width": 200, "height": 200},
|
||||
{"url": "medium.jpg", "preference": 1, "width": 150, "height": 150},
|
||||
]
|
||||
|
||||
result = get_thumbnail(thumbnails)
|
||||
assert result == {"url": "large.jpg", "preference": 1, "width": 200, "height": 200}
|
||||
|
||||
def test_handles_missing_attributes(self):
|
||||
"""Test that thumbnails with missing attributes are handled correctly."""
|
||||
from app.library.Utils import get_thumbnail
|
||||
|
||||
thumbnails = [
|
||||
{"url": "no_pref.jpg", "width": 100},
|
||||
{"url": "with_pref.jpg", "preference": 5, "width": 50},
|
||||
]
|
||||
|
||||
result = get_thumbnail(thumbnails)
|
||||
assert result["url"] == "with_pref.jpg"
|
||||
|
||||
def test_returns_first_when_all_equal(self):
|
||||
"""Test that any thumbnail is returned when all attributes are equal."""
|
||||
from app.library.Utils import get_thumbnail
|
||||
|
||||
thumbnails = [
|
||||
{"url": "first.jpg"},
|
||||
{"url": "second.jpg"},
|
||||
]
|
||||
|
||||
result = get_thumbnail(thumbnails)
|
||||
assert result is not None
|
||||
assert result["url"] in ["first.jpg", "second.jpg"]
|
||||
|
||||
|
||||
class TestGetExtras:
|
||||
def test_returns_empty_dict_for_none(self):
|
||||
"""Test that empty dict is returned for None input."""
|
||||
from app.library.Utils import get_extras
|
||||
|
||||
assert get_extras(None) == {}
|
||||
|
||||
def test_returns_empty_dict_for_non_dict(self):
|
||||
"""Test that empty dict is returned for non-dict input."""
|
||||
from app.library.Utils import get_extras
|
||||
|
||||
assert get_extras("not a dict") == {}
|
||||
assert get_extras([]) == {}
|
||||
|
||||
def test_extracts_video_information(self):
|
||||
"""Test extracting information from a video entry."""
|
||||
from app.library.Utils import get_extras
|
||||
|
||||
entry = {
|
||||
"id": "test123",
|
||||
"title": "Test Video",
|
||||
"uploader": "Test Uploader",
|
||||
"channel": "Test Channel",
|
||||
"thumbnails": [{"url": "thumb.jpg", "preference": 1}],
|
||||
"duration": 120,
|
||||
}
|
||||
|
||||
result = get_extras(entry, kind="video")
|
||||
|
||||
assert result["uploader"] == "Test Uploader"
|
||||
assert result["channel"] == "Test Channel"
|
||||
assert result["thumbnail"] == "thumb.jpg"
|
||||
assert result["duration"] == 120
|
||||
assert result["is_premiere"] is False
|
||||
|
||||
def test_extracts_playlist_information(self):
|
||||
"""Test extracting information from a playlist entry."""
|
||||
from app.library.Utils import get_extras
|
||||
|
||||
entry = {
|
||||
"id": "playlist123",
|
||||
"title": "Test Playlist",
|
||||
"uploader": "Playlist Owner",
|
||||
"uploader_id": "owner123",
|
||||
}
|
||||
|
||||
result = get_extras(entry, kind="playlist")
|
||||
|
||||
assert result["playlist_id"] == "playlist123"
|
||||
assert result["playlist_title"] == "Test Playlist"
|
||||
assert result["playlist_uploader"] == "Playlist Owner"
|
||||
assert result["playlist_uploader_id"] == "owner123"
|
||||
|
||||
def test_handles_release_timestamp(self):
|
||||
"""Test handling of release_timestamp for upcoming content."""
|
||||
from app.library.Utils import get_extras
|
||||
|
||||
entry = {
|
||||
"release_timestamp": 1234567890,
|
||||
}
|
||||
|
||||
result = get_extras(entry)
|
||||
|
||||
assert "release_in" in result
|
||||
assert result["release_in"] == "Fri, 13 Feb 2009 23:31:30 GMT"
|
||||
|
||||
def test_handles_upcoming_live_stream(self):
|
||||
"""Test handling of upcoming live stream."""
|
||||
from app.library.Utils import get_extras
|
||||
|
||||
entry = {
|
||||
"release_timestamp": 1234567890,
|
||||
"live_status": "is_upcoming",
|
||||
}
|
||||
|
||||
result = get_extras(entry)
|
||||
|
||||
assert result["is_live"] == 1234567890
|
||||
assert "release_in" in result
|
||||
|
||||
def test_handles_premiere_flag(self):
|
||||
"""Test handling of is_premiere flag."""
|
||||
from app.library.Utils import get_extras
|
||||
|
||||
entry = {
|
||||
"is_premiere": True,
|
||||
}
|
||||
|
||||
result = get_extras(entry)
|
||||
assert result["is_premiere"] is True
|
||||
|
||||
entry2 = {"is_premiere": False}
|
||||
result2 = get_extras(entry2)
|
||||
assert result2["is_premiere"] is False
|
||||
|
||||
def test_youtube_fallback_thumbnail(self):
|
||||
"""Test fallback thumbnail generation for YouTube videos."""
|
||||
from app.library.Utils import get_extras
|
||||
|
||||
entry = {
|
||||
"id": "dQw4w9WgXcQ",
|
||||
"ie_key": "Youtube",
|
||||
}
|
||||
|
||||
result = get_extras(entry)
|
||||
|
||||
assert result["thumbnail"] == "https://img.youtube.com/vi/dQw4w9WgXcQ/maxresdefault.jpg"
|
||||
|
||||
def test_thumbnail_string_fallback(self):
|
||||
"""Test fallback to thumbnail string when thumbnails list not available."""
|
||||
from app.library.Utils import get_extras
|
||||
|
||||
entry = {
|
||||
"thumbnail": "https://example.com/thumb.jpg",
|
||||
}
|
||||
|
||||
result = get_extras(entry)
|
||||
|
||||
assert result["thumbnail"] == "https://example.com/thumb.jpg"
|
||||
|
|
|
|||
Loading…
Reference in a new issue