diff --git a/app/library/Download.py b/app/library/Download.py index 29474493..ff9c790b 100644 --- a/app/library/Download.py +++ b/app/library/Download.py @@ -4,18 +4,20 @@ import logging import multiprocessing import os import re +import signal import time from datetime import UTC, datetime from email.utils import formatdate from pathlib import Path -import yt_dlp +import yt_dlp.utils from .config import Config from .Events import EventBus, Events from .ffprobe import ffprobe from .ItemDTO import ItemDTO from .Utils import delete_dir, extract_info, extract_ytdlp_logs, load_cookies +from .ytdlp import YTDLP from .YTDLPOpts import YTDLPOpts @@ -264,10 +266,17 @@ class Download: params["logger"] = NestedLogger(self.logger) - cls = yt_dlp.YoutubeDL(params=params) + cls = YTDLP(params=params) self.started_time = int(time.time()) + def mark_cancelled(*_): + cls._interrupted = True + cls.to_screen("[info] Interrupt received, exiting cleanly...") + raise SystemExit(130) # noqa: TRY301 + + signal.signal(signal.SIGUSR1, mark_cancelled) + if isinstance(self.info_dict, dict) and len(self.info_dict) > 1: self.logger.debug(f"Downloading '{self.info.url}' using pre-info.") cls.process_ie_result( @@ -275,10 +284,10 @@ class Download: download=True, extra_info={k: v for k, v in self.info.extras.items() if k not in self.info_dict}, ) - ret = cls._download_retcode + ret: int = cls._download_retcode else: self.logger.debug(f"Downloading using url: {self.info.url}") - ret = cls.download(url_list=[self.info.url]) + ret: int = cls.download(url_list=[self.info.url]) self.status_queue.put({"id": self.id, "status": "finished" if ret == 0 else "error"}) except yt_dlp.utils.ExistingVideoReached as exc: @@ -333,7 +342,7 @@ class Download: return False self.cancel_in_progress = True - procId = self.proc.ident + procId: int | None = self.proc.ident self.logger.info(f"Closing PID='{procId}' download process.") @@ -387,7 +396,10 @@ class Download: try: self.logger.info(f"Killing download process: '{self.proc.ident}'.") - self.proc.kill() + if self.proc.pid and "posix" == os.name: + os.kill(self.proc.pid, signal.SIGUSR1) + else: + self.proc.kill() return True except Exception as e: self.logger.error(f"Failed to kill process: '{self.proc.ident}'. {e}") @@ -398,9 +410,9 @@ class Download: if self.temp_keep is True or not self.temp_path: return - if "finished" != self.info.status and self.is_live: + if "finished" != self.info.status and self.info.downloaded_bytes > 0: self.logger.warning( - f"Keeping live temp folder '{self.temp_path}', as the reported status is not finished '{self.info.status}'." + f"Keeping temp folder '{self.temp_path}', as the reported status is not finished '{self.info.status}'." ) return diff --git a/app/library/DownloadQueue.py b/app/library/DownloadQueue.py index fb1bf668..ff4f4fd0 100644 --- a/app/library/DownloadQueue.py +++ b/app/library/DownloadQueue.py @@ -10,7 +10,7 @@ from pathlib import Path from sqlite3 import Connection from typing import TYPE_CHECKING -import yt_dlp +import yt_dlp.utils from aiohttp import web from .ag_utils import ag diff --git a/app/library/Utils.py b/app/library/Utils.py index 18e90ee3..8f9a8cdc 100644 --- a/app/library/Utils.py +++ b/app/library/Utils.py @@ -15,11 +15,11 @@ from http.cookiejar import MozillaCookieJar from pathlib import Path from typing import TypeVar -import yt_dlp from Crypto.Cipher import AES from yt_dlp.utils import age_restricted, match_str from .LogWrapper import LogWrapper +from .ytdlp import YTDLP LOG: logging.Logger = logging.getLogger("Utils") @@ -47,7 +47,7 @@ REMOVE_KEYS: list = [ }, ] -YTDLP_INFO_CLS: yt_dlp.YoutubeDL = None +YTDLP_INFO_CLS: YTDLP = None ALLOWED_SUBS_EXTENSIONS: tuple[str] = (".srt", ".vtt", ".ass") @@ -189,7 +189,7 @@ def extract_info( if no_archive and "download_archive" in params: del params["download_archive"] - data = yt_dlp.YoutubeDL(params=params).extract_info(url, download=False) + data = YTDLP(params=params).extract_info(url, download=False) if data and follow_redirect and "_type" in data and "url" == data["_type"]: return extract_info( @@ -208,7 +208,7 @@ def extract_info( if not data["is_premiere"]: data["is_premiere"] = "video" == data.get("media_type") and "is_upcoming" == data.get("live_status") - return yt_dlp.YoutubeDL.sanitize_info(data) if sanitize_info else data + return YTDLP.sanitize_info(data) if sanitize_info else data def merge_dict(source: dict, destination: dict) -> dict: @@ -1103,7 +1103,7 @@ def get_archive_id(url: str) -> tuple[bool, dict[str | None, str | None, str | N } if YTDLP_INFO_CLS is None: - YTDLP_INFO_CLS = yt_dlp.YoutubeDL( + YTDLP_INFO_CLS = YTDLP( params={ "color": "no_color", "extract_flat": True, diff --git a/app/library/ytdlp.py b/app/library/ytdlp.py new file mode 100644 index 00000000..8c5cc719 --- /dev/null +++ b/app/library/ytdlp.py @@ -0,0 +1,12 @@ +import yt_dlp + + +class YTDLP(yt_dlp.YoutubeDL): + _interrupted = False + + def _delete_downloaded_files(self, *args, **kwargs): + if self._interrupted: + self.to_screen("[info] Cancelled — skipping temp cleanup.") + return None + + return super()._delete_downloaded_files(*args, **kwargs)