ytptube/app/library/Download.py

723 lines
28 KiB
Python

import asyncio
import hashlib
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
from typing import Any
import yt_dlp.utils
from .config import Config
from .Events import EventBus, Events
from .ffprobe import ffprobe
from .ItemDTO import ItemDTO
from .Utils import create_cookies_file, delete_dir, extract_info, extract_ytdlp_logs
from .ytdlp import YTDLP
class Terminator:
pass
class NestedLogger:
debug_messages: list[str] = ["[debug] ", "[download] "]
def __init__(self, logger: logging.Logger) -> None:
self.logger: logging.Logger = logger
def debug(self, msg: str) -> None:
levelno: int = logging.DEBUG if any(msg.startswith(x) for x in self.debug_messages) else logging.INFO
self.logger.log(level=levelno, msg=re.sub(r"^\[(debug|info)\] ", "", msg, flags=re.IGNORECASE))
def error(self, msg) -> None:
self.logger.error(msg)
def warning(self, msg) -> None:
self.logger.warning(msg)
class Download:
"""
Download task.
"""
temp_path: str = None
update_task = None
cancel_in_progress: bool = False
final_update = False
bad_live_options: list = [
"concurrent_fragment_downloads",
"fragment_retries",
"skip_unavailable_fragments",
]
"Bad yt-dlp options which are known to cause issues with live stream."
_ytdlp_fields: tuple = (
"tmpfilename",
"filename",
"status",
"msg",
"total_bytes",
"total_bytes_estimate",
"downloaded_bytes",
"speed",
"eta",
)
"Fields to be extracted from yt-dlp progress hook."
def __init__(self, info: ItemDTO, info_dict: dict = None, logs: list | None = None):
"""
Initialize download task.
Args:
info (ItemDTO): ItemDTO object.
info_dict (dict): yt-dlp metadata dict.
logs (list): Logs from yt-dlp.
"""
config: Config = Config.get_instance()
self.download_dir: str = info.download_dir
"Download directory."
self.temp_dir: str | None = info.temp_dir
"Temporary directory."
self.template: str | None = info.template
"Filename template."
self.template_chapter: str | None = info.template_chapter
"Chapter filename template."
self.download_info_expires = int(config.download_info_expires)
"Time in seconds before the download info is considered expired."
self.info: ItemDTO = info
"ItemDTO object."
self.id: str = info._id
"Download ID."
self.debug = bool(config.debug)
"Debug mode."
self.debug_ytdl = bool(config.ytdlp_debug)
"Debug mode for yt-dlp."
self.cancelled = False
"Download cancelled."
self.tmpfilename = None
"Temporary filename."
self.status_queue = None
"Status queue."
self.proc = None
"yt-dlp process."
self._notify: EventBus = EventBus.get_instance()
"Event bus instance."
self.max_workers = int(config.max_workers)
"Maximum number of concurrent downloads."
self.temp_keep = bool(config.temp_keep)
"Keep temp folder after download."
self.temp_disabled = bool(config.temp_disabled)
"Disable the temporary files feature."
self.is_live: bool = bool(info.is_live) or info.live_in is not None
"Is the download a live stream."
self.info_dict: dict = info_dict
"yt-dlp metadata dict."
self.logger: logging.Logger = logging.getLogger(f"Download.{info.id if info.id else info._id}")
"Logger for the download task."
self.started_time = 0
"Time when the download started."
self.queue_time: datetime = datetime.now(tz=UTC)
"Time when the download was queued."
self.logs: list[str] = logs if logs else []
"Logs from yt-dlp."
def _progress_hook(self, data: dict) -> None:
if self.debug:
try:
d_safe: dict = {
"status": data.get("status"),
"filename": data.get("filename"),
"info_dict": {
k: v
for k, v in data.get("info_dict", {}).items()
if k not in ["formats", "thumbnails", "description", "tags", "_format_sort_fields"]
and not isinstance(v, (type(None).__bases__[0], type(lambda: None))) # Skip non-serializable
},
}
self.logger.debug(f"PG Hook: {d_safe}")
except Exception as e:
self.logger.debug(f"PG Hook: Error creating debug info: {e}")
self.status_queue.put(
{
"id": self.id,
"action": "progress",
**{k: v for k, v in data.items() if k in self._ytdlp_fields},
}
)
def _postprocessor_hook(self, data: dict) -> None:
if self.debug:
try:
d_safe: dict = {
"postprocessor": data.get("postprocessor"),
"status": data.get("status"),
"info_dict": {
k: v
for k, v in data.get("info_dict", {}).items()
if k not in ["formats", "thumbnails", "description", "tags", "_format_sort_fields"]
and not isinstance(v, (type(None).__bases__[0], type(lambda: None))) # Skip non-serializable
},
}
self.logger.debug(f"PP Hook: {d_safe}")
except Exception as e:
self.logger.debug(f"PP Hook: Error creating debug info: {e}")
if "MoveFiles" == data.get("postprocessor") and "finished" == data.get("status"):
if "__finaldir" in data.get("info_dict", {}) and "filepath" in data.get("info_dict", {}):
filename = str(Path(data["info_dict"]["__finaldir"]) / Path(data["info_dict"]["filepath"]).name)
else:
filename: str | None = data.get("info_dict", {}).get("filepath", data.get("filename"))
self.logger.debug(f"Final filename: '{filename}'.")
self.status_queue.put({"id": self.id, "action": "moved", "status": "finished", "final_name": filename})
return
dataDict = {k: v for k, v in data.items() if k in self._ytdlp_fields}
self.status_queue.put({"id": self.id, "action": "postprocessing", **dataDict, "status": "postprocessing"})
def post_hooks(self, filename: str | None = None) -> None:
if not filename:
return
self.status_queue.put({"id": self.id, "filename": filename})
def _download(self) -> None:
if not self._notify:
self._notify = EventBus.get_instance()
cookie_file = None
try:
params: dict = (
self.info.get_ytdlp_opts()
.add(
config={
"color": "no_color",
"paths": {
"home": str(self.download_dir),
"temp": str(self.temp_path) if self.temp_path else self.download_dir,
},
"outtmpl": {
"default": self.template,
"chapter": self.template_chapter,
},
"noprogress": True,
"ignoreerrors": False,
},
from_user=False,
)
.get_all()
)
params.update(
{
"progress_hooks": [self._progress_hook],
"postprocessor_hooks": [self._postprocessor_hook],
"post_hooks": [self.post_hooks],
}
)
if self.debug_ytdl:
params["verbose"] = True
params["noprogress"] = False
if self.info.cookies:
try:
cookie_file: Path = (
Path(self.temp_path if self.temp_path else self.temp_dir) / f"cookie_{self.info._id}.txt"
)
self.logger.debug(
f"Creating cookie file for '{self.info.id}: {self.info.title}' - '{cookie_file}'."
)
params["cookiefile"] = str(create_cookies_file(self.info.cookies, cookie_file).as_posix())
except Exception as e:
err_msg: str = f"Failed to create cookie file for '{self.info.id}: {self.info.title}'. '{e!s}'."
self.logger.error(err_msg)
raise ValueError(err_msg) from e
# Safe-guard in-case downloading take too long and the info expires.
if self.info_dict and isinstance(self.info_dict, dict) and self.download_info_expires > 0:
_ts: int | None = self.info_dict.get("epoch", self.info_dict.get("timestamp", None))
_ts = datetime.fromtimestamp(_ts, tz=UTC) if _ts else None
if not _ts or (datetime.now(tz=UTC) - _ts).total_seconds() > self.download_info_expires:
self.info_dict = None
self.logger.warning(f"Info for '{self.info.url}' has expired, re-extracting info.")
if not self.info_dict or not isinstance(self.info_dict, dict):
self.logger.info(f"Extracting info for '{self.info.url}'.")
self.logs = []
ie_params: dict = params.copy()
ie_params["callback"] = {
"func": lambda _, msg: self.logs.append(msg),
"level": logging.WARNING,
"name": "callback-logger",
}
info: dict = extract_info(
config=params,
url=self.info.url,
debug=self.debug,
no_archive=not params.get("download_archive", False),
follow_redirect=True,
)
if info:
self.info_dict = info
if self.is_live:
deletedOpts: list = []
for opt in self.bad_live_options:
if opt in params:
params.pop(opt, None)
deletedOpts.append(opt)
if len(deletedOpts) > 0:
self.logger.warning(
f"Live stream detected for '{self.info.title}', The following opts '{deletedOpts=}' have been deleted."
)
if isinstance(self.info_dict, dict) and len(self.info_dict.get("formats", [])) < 1:
msg: str = f"Failed to extract any formats for '{self.info.url}'."
if filtered_logs := extract_ytdlp_logs(self.logs):
msg += " " + ", ".join(filtered_logs)
raise ValueError(msg) # noqa: TRY301
self.logger.info(
f'Task {self.info.name()}, preset="{self.info.preset}", cookies="{bool(params.get("cookiefile"))}" started.'
)
if self.debug:
self.logger.debug(f"Params before passing to yt-dlp. {params}")
params["logger"] = NestedLogger(self.logger)
if "continuedl" in params and params["continuedl"] is False:
self.delete_temp(by_pass=True)
cls = YTDLP(params=params)
self.started_time = int(time.time())
if "posix" == os.name:
def mark_cancelled(*_) -> None:
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.")
_dct: dict = self.info_dict.copy()
if isinstance(self.info.extras, dict) and len(self.info.extras) > 0:
_dct.update({k: v for k, v in self.info.extras.items() if k not in _dct or not _dct.get(k)})
cls.process_ie_result(ie_result=_dct, download=True)
ret: int = cls._download_retcode
else:
self.logger.debug(f"Downloading using url: {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:
self.logger.error(exc)
self.status_queue.put({"id": self.id, "status": "skip", "msg": "Item has already been downloaded."})
except Exception as exc:
self.logger.exception(exc)
self.logger.error(exc)
self.status_queue.put({"id": self.id, "status": "error", "msg": str(exc), "error": str(exc)})
finally:
self.status_queue.put(Terminator())
if cookie_file and cookie_file.exists():
try:
cookie_file.unlink()
self.logger.debug(f"Deleted cookie file: {cookie_file}")
except Exception as e:
self.logger.error(f"Failed to delete cookie file: {cookie_file}. {e}")
self.logger.info(
f'Task {self.info.name()} preset="{self.info.preset}" cookies="{bool(params.get("cookiefile"))}" completed.'
)
async def start(self) -> None:
self.status_queue: multiprocessing.Queue[Any] = Config.get_manager().Queue()
# Create temp dir for each download.
if not self.temp_disabled:
self.temp_path = Path(self.temp_dir) / hashlib.shake_256(f"D-{self.info.id}".encode()).hexdigest(5)
if not self.temp_path.exists():
self.temp_path.mkdir(parents=True, exist_ok=True)
self.info.temp_path = str(self.temp_path)
self.proc = multiprocessing.Process(name=f"download-{self.id}", target=self._download)
self.proc.start()
self.info.status = "preparing"
self._notify.emit(Events.ITEM_UPDATED, data=self.info)
asyncio.create_task(self.progress_update(), name=f"update-{self.id}")
ret = await asyncio.get_running_loop().run_in_executor(None, self.proc.join)
if self.final_update or self.cancelled:
if self.cancelled:
self.info.status = "cancelled"
return ret
self.status_queue.put(Terminator())
self.logger.debug("Draining status queue.")
try:
drain_count: int = 50 + (self.status_queue.qsize() if hasattr(self.status_queue, "qsize") else 5)
except Exception:
drain_count = 55
for i in range(drain_count):
try:
if self.debug:
self.logger.debug(f"(50/{i}) Draining the status queue...")
if self.final_update:
if self.debug:
self.logger.debug("(50/{i}) Draining stopped. Final update received.")
break
next_status = self.status_queue.get(timeout=0.1)
if next_status is None or isinstance(next_status, Terminator):
continue
await self._process_status_update(next_status)
except Exception: # noqa: S112
continue
return ret
def started(self) -> bool:
return self.proc is not None
def cancel(self) -> bool:
"""
Cancel the download task.
For live streams and unresponsive processes, this will force-kill the process.
"""
if not self.started():
return False
self.cancelled = True
# Actively kill the process instead of just setting a flag
return self.kill()
async def close(self) -> bool:
"""
Close download process.
This method MUST be called to clear resources.
"""
if not self.started() or self.cancel_in_progress:
return False
self.cancel_in_progress = True
procId: int | None = self.proc.ident
if not procId:
if self.proc:
self.proc.close()
self.proc = None
self.logger.warning("Attempted to close download process, but it is not running.")
return False
self.logger.info(f"Closing PID='{procId}' download process.")
try:
try:
if "update_task" in self.__dict__ and self.update_task.done() is False:
self.update_task.cancel()
except Exception as e:
self.logger.error(f"Failed to close status queue: '{procId}'. {e}")
self.kill()
loop = asyncio.get_running_loop()
if self.proc.is_alive():
self.logger.debug(f"Waiting for PID='{procId}' to close.")
await loop.run_in_executor(None, self.proc.join)
self.logger.debug(f"PID='{procId}' closed.")
if self.status_queue:
self.logger.debug(f"Closing status queue for PID='{procId}'.")
self.status_queue.put(Terminator())
self.logger.debug(f"Closed status queue for PID='{procId}'.")
if self.proc:
self.proc.close()
self.proc = None
self.delete_temp()
self.logger.debug(f"Closed PID='{procId}' download process.")
return True
except Exception as e:
self.logger.error(f"Failed to close process: '{procId}'. {e}")
self.logger.exception(e)
return False
def running(self) -> bool:
try:
return self.proc is not None and self.proc.is_alive()
except ValueError:
return False
def is_cancelled(self) -> bool:
return self.cancelled
def kill(self) -> bool:
"""
Kill the download process.
Attempts graceful termination via signal first, then force-kills if necessary.
For live streams that may not respond to signals, uses force termination.
"""
if not self.running():
return False
procId: int | None = self.proc.ident
try:
self.logger.info(f"Killing download process: PID={self.proc.pid}, ident={procId}.")
# First, try graceful termination with SIGUSR1 signal (if on POSIX)
if self.proc.pid and "posix" == os.name:
try:
self.logger.debug(f"Sending SIGUSR1 signal to PID={self.proc.pid}.")
os.kill(self.proc.pid, signal.SIGUSR1)
# Wait briefly for graceful shutdown (1 second timeout for live streams)
start_time = time.time()
timeout = 2 if self.is_live else 5 # Live streams get shorter timeout
while self.proc.is_alive() and (time.time() - start_time) < timeout:
time.sleep(0.1)
if self.proc.is_alive():
self.logger.warning(
f"Process PID={self.proc.pid} did not respond to SIGUSR1 "
f"({'live stream' if self.is_live else 'regular download'}), "
f"forcing termination."
)
else:
self.logger.debug(f"Process PID={self.proc.pid} terminated gracefully.")
return True
except (OSError, AttributeError) as e:
self.logger.debug(f"Failed to send SIGUSR1 signal: {e}")
# Force terminate if still running (for live streams or unresponsive processes)
if self.proc.is_alive():
self.logger.info(f"Force-terminating process PID={self.proc.pid}.")
self.proc.terminate()
# Give it a moment to terminate
start_time = time.time()
timeout = 1 if self.is_live else 2
while self.proc.is_alive() and (time.time() - start_time) < timeout:
time.sleep(0.05)
# If still alive, use hard kill
if self.proc.is_alive():
self.logger.warning(
f"Process PID={self.proc.pid} did not respond to terminate(), killing forcefully."
)
self.proc.kill()
# Give it final moment
start_time = time.time()
while self.proc.is_alive() and (time.time() - start_time) < 1:
time.sleep(0.05)
self.logger.info(f"Process PID={self.proc.pid} killed.")
return True
except Exception as e:
self.logger.error(f"Failed to kill process PID={self.proc.pid}, ident={procId}. {e}")
self.logger.exception(e)
return False
def delete_temp(self, by_pass: bool = False) -> None:
if self.temp_disabled or self.temp_keep is True or not self.temp_path:
return
if (
not by_pass
and "finished" != self.info.status
and self.info.downloaded_bytes
and self.info.downloaded_bytes > 0
):
self.logger.warning(f"Keeping temp folder '{self.temp_path}'. {self.info.status=}.")
return
tmp_dir = Path(self.temp_path)
if not tmp_dir.exists():
return
if str(tmp_dir) == str(self.temp_dir):
self.logger.warning(f"Refusing to delete video temp folder '{self.temp_path}' as it's temp root.")
return
status: bool = delete_dir(tmp_dir)
if by_pass:
tmp_dir.mkdir(parents=True, exist_ok=True)
self.logger.info(f"Temp folder '{self.temp_path}' emptied.")
else:
self.logger.info(f"Temp folder '{self.temp_path}' deletion is {'success' if status else 'failed'}.")
async def _process_status_update(self, status) -> None:
if status.get("id") != self.id or len(status) < 2:
self.logger.warning(f"Received invalid status update. {status}")
return
if self.debug:
self.logger.debug(f"Status Update: {self.info._id=} {status=}")
if isinstance(status, str):
self._notify.emit(Events.ITEM_UPDATED, data=self.info)
return
self.tmpfilename: str | None = status.get("tmpfilename")
fl = None
if "final_name" in status:
fl = Path(status.get("final_name"))
try:
self.info.filename = str(fl.relative_to(Path(self.download_dir)))
except ValueError as ve:
self.logger.debug(f"Failed to get relative path for '{fl}' from '{self.download_dir}'. {ve}")
if self.temp_path:
self.info.filename = str(fl.relative_to(Path(self.temp_path)))
else:
self.info.filename = str(fl)
if fl.is_file() and fl.exists():
self.final_update = True
self.logger.debug(f"Final file name: '{fl}'.")
try:
self.info.file_size = fl.stat().st_size
except FileNotFoundError:
self.info.file_size = 0
self.info.status = status.get("status", self.info.status)
self.info.msg = status.get("msg")
if "error" == self.info.status and "error" in status:
self.info.error = status.get("error")
self._notify.emit(
Events.LOG_ERROR,
data=self.info,
title="Download Error",
message=f"'{self.info.title}' failed to download. {self.info.error}",
)
if "downloaded_bytes" in status and status.get("downloaded_bytes", 0) > 0:
self.info.downloaded_bytes = status.get("downloaded_bytes")
total: float | None = status.get("total_bytes") or status.get("total_bytes_estimate")
if total:
try:
self.info.percent = status["downloaded_bytes"] / total * 100
except ZeroDivisionError:
self.info.percent = 0
self.info.total_bytes = total
self.info.speed = status.get("speed")
self.info.eta = status.get("eta")
if "finished" == self.info.status and fl and fl.is_file() and fl.exists():
self.info.file_size = fl.stat().st_size
self.info.datetime = str(formatdate(time.time()))
try:
ff = await ffprobe(str(fl))
self.info.extras["is_video"] = ff.has_video()
self.info.extras["is_audio"] = ff.has_audio()
if (ff.has_video() or ff.has_audio()) and not self.info.extras.get("duration"):
self.info.extras["duration"] = int(float(ff.metadata.get("duration", 0.0)))
except Exception as e:
self.info.extras["is_video"] = True
self.info.extras["is_audio"] = True
self.logger.exception(e)
self.logger.error(f"Failed to run ffprobe. {status.get}. {e}")
if not self.final_update or fl:
self._notify.emit(Events.ITEM_UPDATED, data=self.info)
async def progress_update(self) -> None:
"""
Update status of download task and notify the client.
"""
while True:
try:
self.update_task = asyncio.get_running_loop().run_in_executor(None, self.status_queue.get)
status = await self.update_task
if status is None or isinstance(status, Terminator):
return
await self._process_status_update(status)
except (asyncio.CancelledError, OSError, FileNotFoundError):
return
def is_stale(self) -> bool:
"""
Check if the download task is stale.
A download task is considered stale if it has not been updated for a certain period of time.
Returns:
bool: True if the download task is stale, False otherwise.
"""
if not self.info.auto_start:
return False
if self.started_time < 1:
self.logger.debug(f"Download task '{self.info.name()}' not started yet.")
return False
if int(time.time()) - self.started_time < 300:
self.logger.debug(
f"Download task '{self.info.title}: {self.info.id}' started for '{int(time.time()) - self.started_time}' seconds."
)
return False
if not self.running():
self.logger.warning(
f"Download task '{self.info.title}: {self.info.id}' not started running for '{int(time.time()) - self.started_time}' seconds."
)
return True
if self.info.status not in ["finished", "error", "cancelled", "downloading", "postprocessing"]:
status: str = self.info.status if self.info.status else "unknown"
self.logger.warning(
f"Download task '{self.info.title}: {self.info.id}' has been stuck in '{status}' state for '{int(time.time()) - self.started_time}' seconds."
)
return True
return False
def __getstate__(self) -> dict[str, Any]:
state: dict[str, Any] = self.__dict__.copy()
# Exclude (unpickleable) keys during pickling, this issue arise mostly on Windows.
excluded_keys: tuple[str] = ("_notify",)
for key in excluded_keys:
if key in state:
state[key] = None
return state