danvics/ytptube
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

refactor: implement re-extraction logic and process-safe handling for non-picklable objects.

Browse source
This commit is contained in:
arabcoders 2026-05-25 20:06:34 +03:00
parent 2e9684de5d
commit 64fdcac0f7
7 changed files with 715 additions and 500 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

31
app/features/ytdlp/extractor.py
View file

@ -17,6 +17,9 @@ from app.library.Singleton import Singleton
LOG: logging.Logger = logging.getLogger("downloads.extractor")
LIVE_REEXTRACT_STATUSES: set[str] = {"is_live", "post_live"}
REEXTRACT_INFO_KEY = "_ytptube_reextract"
def _ytdlp_logger(target: logging.Logger):
return _YTDLPLogger(target)
@ -247,6 +250,23 @@ def _sanitize_config(config: dict[str, Any]) -> dict[str, Any]:
return sanitized
def needs_reextract(info: dict[str, Any]) -> bool:
return bool(info.get("is_live") or info.get("live_status") in LIVE_REEXTRACT_STATUSES)
def _process_safe_info(info: dict[str, Any]) -> dict[str, Any]:
if needs_reextract(info):
info = dict(info)
info[REEXTRACT_INFO_KEY] = True
for key in ("formats", "requested_formats", "requested_downloads", "fragments"):
info.pop(key, None)
if _is_picklable(info):
return info
return YTDLP.sanitize_info(info, remove_private_keys=False)
def extract_info_sync(
config: dict[str, Any],
url: str,
@ -255,6 +275,7 @@ def extract_info_sync(
follow_redirect: bool = False,
sanitize_info: bool = False,
capture_logs: int | None = None,
process_safe: bool = False,
**kwargs,
) -> tuple[dict[str, Any] | None, list[str]]:
"""
@ -268,6 +289,7 @@ def extract_info_sync(
follow_redirect: Follow URL redirects
sanitize_info: Sanitize the extracted information
capture_logs: If provided (e.g., logging.WARNING), capture logs at this level.
process_safe: Strip non-pickleable data for safe inter-process communication.
**kwargs: Additional arguments
Returns:
@ -330,6 +352,7 @@ def extract_info_sync(
follow_redirect=follow_redirect,
sanitize_info=sanitize_info,
capture_logs=capture_logs,
process_safe=process_safe,
captured_logs=captured_logs,
**kwargs,
)
@ -348,6 +371,9 @@ def extract_info_sync(
result = YTDLP.sanitize_info(data, remove_private_keys=True) if sanitize_info else data
if process_safe and isinstance(result, dict):
result = _process_safe_info(result)
return (result, captured_logs)
finally:
logging.Logger.manager.loggerDict.pop(logger_name, None)
@ -404,6 +430,8 @@ async def fetch_info(
loop = asyncio.get_running_loop()
safe_config = _sanitize_config(config)
safe_kwargs = _sanitize_picklable(kwargs)
safe_kwargs.pop("process_safe", None)
timeout = _sleep_timeout(safe_config, extractor_config.timeout, budget_sleep)
try:
@ -422,7 +450,8 @@ async def fetch_info(
follow_redirect=follow_redirect,
sanitize_info=sanitize_info,
capture_logs=capture_logs,
**kwargs,
process_safe=True,
**safe_kwargs,
),
),
timeout=timeout,

40
app/features/ytdlp/tests/test_ytdlp_extractor.py
View file

@ -5,8 +5,10 @@ from unittest.mock import MagicMock, patch
from app.features.ytdlp.extractor import (
ExtractorConfig,
ExtractorPool,
REEXTRACT_INFO_KEY,
_LogCapture,
_get_process_pool_kwargs,
_process_safe_info,
_ytdlp_logger,
extract_info_sync,
)
@ -158,6 +160,44 @@ class TestExtractInfo:
assert (logging.INFO, "[generic_browser] Using remote browser for https://example.com/video") in seen
assert (logging.WARNING, "[generic_browser] Browser fallback warning") in seen
def test_process_safe_live(self) -> None:
data = {
"id": "live-id",
"is_live": True,
"formats": [{"format_id": "dash", "fragments": ({"url": "https://example.test/sq/1"} for _ in range(1))}],
"requested_formats": [{"format_id": "dash"}],
}
result = _process_safe_info(data)
assert result[REEXTRACT_INFO_KEY] is True
assert "formats" not in result
assert "requested_formats" not in result
pickle.dumps(result)
def test_process_safe_post_live(self) -> None:
data = {
"id": "post-live-id",
"live_status": "post_live",
"formats": [{"format_id": "dash", "fragments": ({"url": "https://example.test/sq/1"} for _ in range(1))}],
}
result = _process_safe_info(data)
assert result[REEXTRACT_INFO_KEY] is True
assert "formats" not in result
pickle.dumps(result)
def test_process_safe_lazy(self) -> None:
from yt_dlp.utils import LazyList
data = {"id": "video-id", "formats": LazyList({"format_id": str(i)} for i in range(2))}
result = _process_safe_info(data)
assert result["formats"] == [{"format_id": "0"}, {"format_id": "1"}]
pickle.dumps(result)
class TestYtdlpLogger:
def test_debug_prefix_uses_debug(self) -> None:

11
app/library/downloads/core.py
View file

@ -21,7 +21,7 @@ from app.library.config import Config
from app.library.Events import EventBus, Events
from app.library.Utils import create_cookies_file
from ...features.ytdlp.extractor import extract_info_sync
from ...features.ytdlp.extractor import REEXTRACT_INFO_KEY, extract_info_sync
from .hooks import HookHandlers, NestedLogger
from .process_manager import ProcessManager
from .status_tracker import StatusTracker
@ -147,11 +147,16 @@ class Download:
raise ValueError(err_msg) from e
if self.info_dict and isinstance(self.info_dict, dict):
if self.info_dict.get("extractor_key") in GENERIC_EXTRACTORS and self.info.get_preset().default:
if self.info_dict.get(REEXTRACT_INFO_KEY):
self.logger.info(
f"Info dict for '{self.info.url}' marked for re-extraction, ignoring pre-extracted info."
)
self.info_dict = None
elif self.info_dict.get("extractor_key") in GENERIC_EXTRACTORS and self.info.get_preset().default:
self.logger.debug(f"Removing 'download_archive' for generic extractor. url={self.info.url}")
params.pop("download_archive", None)
if self.download_info_expires > 0:
if self.info_dict and self.download_info_expires > 0:
_ts: int | None = self.info_dict.get("epoch", self.info_dict.get("timestamp", None))
_ts_dt = datetime.fromtimestamp(_ts, tz=UTC) if _ts else None
if not _ts_dt or (datetime.now(tz=UTC) - _ts_dt).total_seconds() > self.download_info_expires:

6
app/library/downloads/video_processor.py
View file

@ -6,6 +6,7 @@ from datetime import UTC, datetime, timedelta
from email.utils import formatdate
from typing import TYPE_CHECKING
from app.features.ytdlp.extractor import REEXTRACT_INFO_KEY, needs_reextract
from app.features.ytdlp.utils import extract_ytdlp_logs, get_extras
from app.library.downloads import Download
from app.library.Events import Events
@ -125,12 +126,13 @@ async def add_video(queue: "DownloadQueue", entry: dict, item: "Item", logs: lis
)
try:
dlInfo: Download = Download(info=dl, info_dict=entry if item.auto_start else None, logs=logs)
_reextract: bool = bool(entry.get(REEXTRACT_INFO_KEY) or needs_reextract(entry))
dlInfo: Download = Download(info=dl, info_dict=entry if item.auto_start and not _reextract else None, logs=logs)
nEvent: str | None = None
nTitle: str | None = None
nMessage: str | None = None
nStore: str = "queue"
hasFormats: bool = len(entry.get("formats", [])) > 0 or entry.get("url")
hasFormats: bool = _reextract or bool(entry.get("formats") or entry.get("url"))
text_logs: str = ""
if filtered_logs := extract_ytdlp_logs(logs):

92
app/tests/test_download.py
View file

@ -1,6 +1,7 @@
import logging
import os
import signal
import time
from multiprocessing.reduction import ForkingPickler
from pathlib import Path
from types import SimpleNamespace
@ -17,6 +18,7 @@ from app.library.downloads.process_manager import ProcessManager
from app.library.downloads.queue_manager import DownloadQueue
from app.library.downloads.status_tracker import StatusTracker
from app.library.downloads.temp_manager import TempManager
from app.library.downloads.video_processor import add_video
from app.library.ItemDTO import ItemDTO
@ -292,7 +294,12 @@ class TestDownloadFlow:
download = Download(
info=make_item(),
info_dict={"id": "test-id", "url": "http://u", "formats": [{"format_id": "18"}]},
info_dict={
"id": "test-id",
"url": "http://u",
"formats": [{"format_id": "18"}],
"epoch": int(time.time()),
},
)
download.status_queue = cast(Any, DummyQueue())
download._hook_handlers = Mock(
@ -326,6 +333,7 @@ class TestDownloadFlow:
queue = cast(DummyQueue, download.status_queue)
assert queue.items[0]["download_skipped"] is True
assert queue.items[1]["status"] == "finished"
assert queue.items[1]["download_skipped"] is True
def test_playlist_extras(self, monkeypatch: pytest.MonkeyPatch) -> None:
@ -1205,6 +1213,88 @@ class TestStatusTracker:
class TestQueueManager:
@staticmethod
def _video_queue() -> Mock:
async def put(item):
return item
queue_manager = Mock()
queue_manager.config = Mock(
download_path="/tmp",
temp_path="/tmp",
output_template="%(title)s.%(ext)s",
output_template_chapter="%(title)s.%(ext)s",
prevent_live_premiere=False,
)
queue_manager.done.get = AsyncMock(side_effect=KeyError)
queue_manager.queue.get = AsyncMock(side_effect=KeyError)
queue_manager.done.put = AsyncMock(side_effect=put)
queue_manager.queue.put = AsyncMock(side_effect=put)
queue_manager.pool.trigger_download = Mock()
queue_manager._notify.emit = Mock()
return queue_manager
@staticmethod
def _video_item() -> SimpleNamespace:
return SimpleNamespace(
extras={},
folder="",
preset="default",
cookies=None,
template=None,
cli=[],
auto_start=True,
)
@pytest.mark.asyncio
async def test_live_reextracts(self, monkeypatch: pytest.MonkeyPatch) -> None:
seen: list[dict | None] = []
def fake_download(*, info, info_dict, logs):
seen.append(info_dict)
return SimpleNamespace(info=info, info_dict=info_dict, logs=logs)
monkeypatch.setattr("app.library.downloads.video_processor.Download", fake_download)
result = await add_video(
queue=self._video_queue(),
item=self._video_item(),
entry={
"id": "live-id",
"title": "Live stream",
"webpage_url": "https://example.test/live",
"is_live": True,
"live_status": "is_live",
"_ytptube_reextract": True,
},
)
assert result == {"status": "ok"}
assert seen == [None]
@pytest.mark.asyncio
async def test_regular_reuses_info(self, monkeypatch: pytest.MonkeyPatch) -> None:
seen: list[dict | None] = []
def fake_download(*, info, info_dict, logs):
seen.append(info_dict)
return SimpleNamespace(info=info, info_dict=info_dict, logs=logs)
monkeypatch.setattr("app.library.downloads.video_processor.Download", fake_download)
entry = {
"id": "video-id",
"title": "Video",
"webpage_url": "https://example.test/video",
"live_status": "not_live",
"formats": [{"format_id": "18"}],
}
result = await add_video(queue=self._video_queue(), item=self._video_item(), entry=entry)
assert result == {"status": "ok"}
assert seen == [entry]
@pytest.mark.asyncio
async def test_cleanup_thumbnails(self, tmp_path: Path) -> None:
from app.library.downloads.monitors import cleanup_thumbnails

27
app/tests/test_extractor_timeout.py
View file

@ -92,3 +92,30 @@ async def test_pool_fallback(monkeypatch: pytest.MonkeyPatch) -> None:
assert loop.calls == [pool.executor, None]
assert seen == [130, 130]
assert not pool.semaphore.locked()
@pytest.mark.asyncio
async def test_pool_process_safe(monkeypatch: pytest.MonkeyPatch) -> None:
pool = _Pool()
loop = _Loop()
expected = ({"id": "ok"}, [])
sync = Mock(return_value=expected)
async def fake_wait_for(*, fut, timeout): # noqa: ARG001
return fut()
monkeypatch.setattr(extractor.ExtractorPool, "get_instance", classmethod(lambda cls: pool))
monkeypatch.setattr(extractor.asyncio, "get_running_loop", lambda: loop)
monkeypatch.setattr(extractor.asyncio, "wait_for", fake_wait_for)
monkeypatch.setattr(extractor, "extract_info_sync", sync)
result = await extractor.fetch_info(
config={},
url="https://example.com",
extractor_config=extractor.ExtractorConfig(concurrency=1, timeout=70),
)
assert result == expected
assert loop.calls == [pool.executor]
assert sync.call_args.kwargs["process_safe"] is True
assert not pool.semaphore.locked()

1008
uv.lock
View file

File diff suppressed because it is too large Load diff