import contextlib import copy import json import os import time from contextlib import contextmanager from pathlib import Path from typing import Callable from typing import Dict from typing import List from typing import Optional import yt_dlp as ytdl from yt_dlp.utils import ExistingVideoReached from yt_dlp.utils import MaxDownloadsReached from yt_dlp.utils import RejectedVideoReached from ytdl_sub.thread.log_entries_downloaded_listener import LogEntriesDownloadedListener from ytdl_sub.utils.exceptions import FileNotDownloadedException from ytdl_sub.utils.logger import Logger class YTDLP: _EXTRACT_ENTRY_NUM_RETRIES: int = 5 _EXTRACT_ENTRY_RETRY_WAIT_SEC: int = 5 logger = Logger.get(name="yt-dlp-downloader") @classmethod @contextmanager def ytdlp_downloader(cls, ytdl_options_overrides: Dict) -> ytdl.YoutubeDL: """ Context manager to interact with yt_dlp. """ cls.logger.debug("ytdl_options: %s", str(ytdl_options_overrides)) with Logger.handle_external_logs(name="yt-dlp"): # Deep copy ytdl_options in case yt-dlp modifies the dict with ytdl.YoutubeDL(copy.deepcopy(ytdl_options_overrides)) as ytdl_downloader: yield ytdl_downloader @classmethod def extract_info(cls, ytdl_options_overrides: Dict, **kwargs) -> Dict: """ Wrapper around yt_dlp.YoutubeDL.YoutubeDL.extract_info All kwargs will passed to the extract_info function. Parameters ---------- ytdl_options_overrides Optional. Dict containing ytdl args to override other predefined ytdl args **kwargs arguments passed directory to YoutubeDL extract_info """ with cls.ytdlp_downloader(ytdl_options_overrides) as ytdlp: return ytdlp.extract_info(**kwargs) @classmethod def extract_info_with_retry( cls, ytdl_options_overrides: Dict, is_downloaded_fn: Optional[Callable[[], bool]] = None, is_thumbnail_downloaded_fn: Optional[Callable[[], bool]] = None, **kwargs, ) -> Dict: """ Wrapper around yt_dlp.YoutubeDL.YoutubeDL.extract_info All kwargs will passed to the extract_info function. This should be used when downloading a single entry. Checks if the entry's video and thumbnail files exist - retry if they do not. Parameters ---------- ytdl_options_overrides Dict containing ytdl args to override other predefined ytdl args is_downloaded_fn Optional. Function to check if the entry is downloaded is_thumbnail_downloaded_fn Optional. Function to check if the entry thumbnail is downloaded **kwargs arguments passed directory to YoutubeDL extract_info Raises ------ FileNotDownloadedException If the entry fails to download """ num_tries = 0 entry_files_exist = False copied_ytdl_options_overrides = copy.deepcopy(ytdl_options_overrides) while not entry_files_exist and num_tries < cls._EXTRACT_ENTRY_NUM_RETRIES: entry_dict = cls.extract_info( ytdl_options_overrides=copied_ytdl_options_overrides, **kwargs ) is_downloaded = is_downloaded_fn is None or is_downloaded_fn() is_thumbnail_downloaded = ( is_thumbnail_downloaded_fn is None or is_thumbnail_downloaded_fn() ) if is_downloaded and is_thumbnail_downloaded: return entry_dict # If the video file is downloaded but the thumbnail is not, then do not download # the video again if is_downloaded and not is_thumbnail_downloaded: copied_ytdl_options_overrides["skip_download"] = True copied_ytdl_options_overrides["writethumbnail"] = True time.sleep(cls._EXTRACT_ENTRY_RETRY_WAIT_SEC) num_tries += 1 # Remove the download archive so it can retry without thinking its already downloaded, # even though it is not if "download_archive" in copied_ytdl_options_overrides: del copied_ytdl_options_overrides["download_archive"] if num_tries < cls._EXTRACT_ENTRY_NUM_RETRIES: cls.logger.debug( "Failed to download entry. Retrying %d / %d", num_tries, cls._EXTRACT_ENTRY_NUM_RETRIES, ) error_dict = {"ytdl_options": ytdl_options_overrides, "kwargs": kwargs} raise FileNotDownloadedException( f"yt-dlp failed to download an entry with these arguments: {error_dict}" ) @classmethod def _get_entry_dicts_from_info_json_files(cls, working_directory: str) -> List[Dict]: """ Parameters ---------- working_directory Directory that info json files are located Returns ------- List of all info.json files read as JSON dicts """ entry_dicts: List[Dict] = [] info_json_paths = [ Path(working_directory) / file_name for file_name in os.listdir(working_directory) if file_name.endswith(".info.json") ] for info_json_path in info_json_paths: with open(info_json_path, "r", encoding="utf-8") as file: entry_dicts.append(json.load(file)) return entry_dicts @classmethod @contextlib.contextmanager def _listen_and_log_downloaded_info_json( cls, working_directory: str, log_prefix: Optional[str] ): """ Context manager that starts a separate thread that listens for new .info.json files, prints their titles as they appear """ if not log_prefix: yield return info_json_listener = LogEntriesDownloadedListener( working_directory=working_directory, log_prefix=log_prefix, ) info_json_listener.start() try: yield finally: info_json_listener.complete = True @classmethod def extract_info_via_info_json( cls, working_directory: str, ytdl_options_overrides: Dict, log_prefix_on_info_json_dl: Optional[str] = None, **kwargs, ) -> List[Dict]: """ Wrapper around yt_dlp.YoutubeDL.YoutubeDL.extract_info with infojson enabled. Entry dicts are extracted via reading all info.json files in the working directory rather than from the output of extract_info. This allows us to catch RejectedVideoReached and ExistingVideoReached exceptions, and simply ignore while still being able to read downloaded entry metadata. Parameters ---------- working_directory Directory that info json files reside in ytdl_options_overrides Dict containing ytdl args to override other predefined ytdl args log_prefix_on_info_json_dl Optional. Spin a new thread to listen for new info.json files. Log f'{log_prefix_on_info_json_dl} {title}' when a new one appears **kwargs arguments passed directory to YoutubeDL extract_info """ try: with cls._listen_and_log_downloaded_info_json( working_directory=working_directory, log_prefix=log_prefix_on_info_json_dl ): _ = cls.extract_info(ytdl_options_overrides=ytdl_options_overrides, **kwargs) except RejectedVideoReached: cls.logger.debug( "RejectedVideoReached, stopping additional downloads " "(Can be disable by setting `ytdl_options.break_on_reject` to False)." ) except ExistingVideoReached: cls.logger.debug( "ExistingVideoReached, stopping additional downloads. " "(Can be disable by setting `ytdl_options.break_on_existing` to False)." ) except MaxDownloadsReached: cls.logger.info("MaxDownloadsReached, stopping additional downloads.") return cls._get_entry_dicts_from_info_json_files(working_directory=working_directory)