diff --git a/src/ytdl_sub/downloaders/downloader.py b/src/ytdl_sub/downloaders/downloader.py index c240358b..481a2164 100644 --- a/src/ytdl_sub/downloaders/downloader.py +++ b/src/ytdl_sub/downloaders/downloader.py @@ -20,8 +20,7 @@ from ytdl_sub.downloaders.generic.validators import MultiUrlValidator from ytdl_sub.downloaders.generic.validators import UrlThumbnailListValidator from ytdl_sub.downloaders.generic.validators import UrlValidator from ytdl_sub.downloaders.ytdl_options_builder import YTDLOptionsBuilder -from ytdl_sub.downloaders.ytdlp import extract_info_via_info_json -from ytdl_sub.downloaders.ytdlp import extract_info_with_retry +from ytdl_sub.downloaders.ytdlp import YTDLP from ytdl_sub.entries.entry import Entry from ytdl_sub.entries.entry_parent import EntryParent from ytdl_sub.entries.variables.kwargs import COMMENTS @@ -266,7 +265,7 @@ class YtDlpDownloader(BaseDownloader[DownloaderOptionsT], ABC): FileHandler.delete(info_json_file) def _extract_entry_info_with_retry(self, entry: Entry) -> Entry: - download_entry_dict = extract_info_with_retry( + download_entry_dict = YTDLP.extract_info_with_retry( ytdl_options_overrides=self.download_ytdl_options, is_downloaded_fn=None if self.is_dry_run else entry.is_downloaded, is_thumbnail_downloaded_fn=None @@ -334,7 +333,7 @@ class YtDlpDownloader(BaseDownloader[DownloaderOptionsT], ABC): url = self.overrides.apply_formatter(collection_url.url) with self._separate_download_archives(): - entry_dicts = extract_info_via_info_json( + entry_dicts = YTDLP.extract_info_via_info_json( working_directory=self.working_directory, ytdl_options_overrides=self.metadata_ytdl_options, log_prefix_on_info_json_dl="Downloading metadata for", diff --git a/src/ytdl_sub/downloaders/ytdlp.py b/src/ytdl_sub/downloaders/ytdlp.py index b3993d82..8b2b74ff 100644 --- a/src/ytdl_sub/downloaders/ytdlp.py +++ b/src/ytdl_sub/downloaders/ytdlp.py @@ -19,199 +19,209 @@ from ytdl_sub.thread.log_entries_downloaded_listener import LogEntriesDownloaded from ytdl_sub.utils.exceptions import FileNotDownloadedException from ytdl_sub.utils.logger import Logger -download_logger = Logger.get(name="yt-dlp-downloader") -_extract_entry_num_retries: int = 5 -_extract_entry_retry_wait_sec: int = 5 +class YTDLP: + _EXTRACT_ENTRY_NUM_RETRIES: int = 5 + _EXTRACT_ENTRY_RETRY_WAIT_SEC: int = 5 + logger = Logger.get(name="yt-dlp-downloader") -@contextmanager -def ytdlp_downloader(ytdl_options_overrides: Dict) -> ytdl.YoutubeDL: - """ - Context manager to interact with yt_dlp. - """ - download_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 + @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. -def extract_info(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) - Parameters - ---------- - ytdl_options_overrides - Optional. Dict containing ytdl args to override other predefined ytdl args - **kwargs - arguments passed directory to YoutubeDL extract_info - """ - with 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. -def extract_info_with_retry( - 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. + 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 - This should be used when downloading a single entry. Checks if the entry's video - and thumbnail files exist - retry if they do not. + Raises + ------ + FileNotDownloadedException + If the entry fails to download + """ + num_tries = 0 + entry_files_exist = False + copied_ytdl_options_overrides = copy.deepcopy(ytdl_options_overrides) - 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 < _extract_entry_num_retries: - entry_dict = 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(_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 < _extract_entry_num_retries: - download_logger.debug( - "Failed to download entry. Retrying %d / %d", - num_tries, - _extract_entry_num_retries, + 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 ) - error_dict = {"ytdl_options": ytdl_options_overrides, "kwargs": kwargs} - raise FileNotDownloadedException( - f"yt-dlp failed to download an entry with these arguments: {error_dict}" - ) + 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 -def _get_entry_dicts_from_info_json_files(working_directory: str) -> List[Dict]: - """ - Parameters - ---------- - working_directory - Directory that info json files are located + # 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 - 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") - ] + time.sleep(cls._EXTRACT_ENTRY_RETRY_WAIT_SEC) + num_tries += 1 - 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)) + # 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"] - return entry_dicts + 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, + ) - -@contextlib.contextmanager -def _listen_and_log_downloaded_info_json(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 - - -def extract_info_via_info_json( - 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 _listen_and_log_downloaded_info_json( - working_directory=working_directory, log_prefix=log_prefix_on_info_json_dl - ): - _ = extract_info(ytdl_options_overrides=ytdl_options_overrides, **kwargs) - except RejectedVideoReached: - download_logger.debug( - "RejectedVideoReached, stopping additional downloads " - "(Can be disable by setting `ytdl_options.break_on_reject` to False)." + error_dict = {"ytdl_options": ytdl_options_overrides, "kwargs": kwargs} + raise FileNotDownloadedException( + f"yt-dlp failed to download an entry with these arguments: {error_dict}" ) - except ExistingVideoReached: - download_logger.debug( - "ExistingVideoReached, stopping additional downloads. " - "(Can be disable by setting `ytdl_options.break_on_existing` to False)." - ) - except MaxDownloadsReached: - download_logger.info("MaxDownloadsReached, stopping additional downloads.") - return _get_entry_dicts_from_info_json_files(working_directory=working_directory) + @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) diff --git a/tests/e2e/bandcamp/test_bandcamp.py b/tests/e2e/bandcamp/test_bandcamp.py index 8d0703d0..b3d86034 100644 --- a/tests/e2e/bandcamp/test_bandcamp.py +++ b/tests/e2e/bandcamp/test_bandcamp.py @@ -4,6 +4,7 @@ from expected_download import assert_expected_downloads from expected_transaction_log import assert_transaction_log_matches import ytdl_sub.downloaders.downloader +from ytdl_sub.downloaders.ytdlp import YTDLP from ytdl_sub.subscriptions.subscription import Subscription @@ -65,7 +66,7 @@ class TestBandcamp: # Ensure another invocation will hit ExistingVideoReached if not dry_run: with assert_logs( - logger=ytdl_sub.downloaders.downloader.download_logger, + logger=YTDLP.logger, expected_message="ExistingVideoReached, stopping additional downloads", log_level="debug", ): diff --git a/tests/e2e/plugins/test_date_range.py b/tests/e2e/plugins/test_date_range.py index a2a5de42..2e2f7cf2 100644 --- a/tests/e2e/plugins/test_date_range.py +++ b/tests/e2e/plugins/test_date_range.py @@ -7,6 +7,7 @@ from expected_download import assert_expected_downloads from expected_transaction_log import assert_transaction_log_matches import ytdl_sub.downloaders.downloader +from ytdl_sub.downloaders.ytdlp import YTDLP from ytdl_sub.subscriptions.subscription import Subscription @@ -68,7 +69,7 @@ class TestDateRange: if not dry_run: # try downloading again, ensure nothing more was downloaded with assert_logs( - logger=ytdl_sub.downloaders.downloader.download_logger, + logger=YTDLP.logger, expected_message="ExistingVideoReached, stopping additional downloads", log_level="debug", ): @@ -139,7 +140,7 @@ class TestDateRange: # First, download recent vids. Always download since we want to test dry-run # on the rolling recent portion. with assert_logs( - logger=ytdl_sub.downloaders.downloader.download_logger, + logger=YTDLP.logger, expected_message="RejectedVideoReached, stopping additional downloads", log_level="debug", ): @@ -159,7 +160,7 @@ class TestDateRange: # Then, download the rolling recent vids subscription. This should remove one of the # two videos with assert_logs( - logger=ytdl_sub.downloaders.downloader.download_logger, + logger=YTDLP.logger, expected_message="ExistingVideoReached, stopping additional downloads", log_level="debug", ): @@ -186,7 +187,7 @@ class TestDateRange: # existing if not dry_run: with assert_logs( - logger=ytdl_sub.downloaders.downloader.download_logger, + logger=YTDLP.logger, expected_message="ExistingVideoReached, stopping additional downloads", log_level="debug", ): diff --git a/tests/e2e/youtube/test_playlist.py b/tests/e2e/youtube/test_playlist.py index ac81a38e..5ec58cee 100644 --- a/tests/e2e/youtube/test_playlist.py +++ b/tests/e2e/youtube/test_playlist.py @@ -5,6 +5,7 @@ from expected_download import assert_expected_downloads from expected_transaction_log import assert_transaction_log_matches import ytdl_sub.downloaders.downloader +from ytdl_sub.downloaders.ytdlp import YTDLP from ytdl_sub.subscriptions.subscription import Subscription @@ -83,7 +84,7 @@ class TestPlaylist: # Ensure another invocation will hit ExistingVideoReached if not dry_run: with assert_logs( - logger=ytdl_sub.downloaders.downloader.download_logger, + logger=YTDLP.logger, expected_message="ExistingVideoReached, stopping additional downloads", log_level="debug", ): @@ -131,7 +132,7 @@ class TestPlaylist: if not dry_run: # Ensure another invocation will hit ExistingVideoReached with assert_logs( - logger=ytdl_sub.downloaders.downloader.download_logger, + logger=YTDLP.logger, expected_message="ExistingVideoReached, stopping additional downloads", log_level="debug", ): diff --git a/tests/expected_download.py b/tests/expected_download.py index ab53baed..5ddd2f7b 100644 --- a/tests/expected_download.py +++ b/tests/expected_download.py @@ -1,6 +1,5 @@ import json import os.path -import sys from dataclasses import dataclass from pathlib import Path from typing import List diff --git a/tests/unit/prebuilt_presets/conftest.py b/tests/unit/prebuilt_presets/conftest.py index a78af19f..eea16c51 100644 --- a/tests/unit/prebuilt_presets/conftest.py +++ b/tests/unit/prebuilt_presets/conftest.py @@ -12,6 +12,7 @@ from resources import copy_file_fixture from ytdl_sub.config.config_file import ConfigFile from ytdl_sub.downloaders.downloader import YtDlpDownloader +from ytdl_sub.downloaders.ytdlp import YTDLP from ytdl_sub.entries.variables.kwargs import DESCRIPTION from ytdl_sub.entries.variables.kwargs import EPOCH from ytdl_sub.entries.variables.kwargs import EXT @@ -100,23 +101,22 @@ def mock_entry_dict_factory(mock_downloaded_file_path) -> Callable: @pytest.fixture def mock_download_collection_thumbnail(mock_downloaded_file_path): - def _mock_download_thumbnail(output_path: str) -> bool: - # mock_file_factory(file_name=output_path.split("/")[-1]) - output_name = os.path.basename(output_path) + def _mock_download_and_convert_url_thumbnail( + thumbnail_url: str, output_thumbnail_path: str + ) -> bool: + _ = thumbnail_url + output_name = os.path.basename(output_thumbnail_path) if "poster" in output_name or "show" in output_name: - copy_file_fixture(fixture_name="poster.jpg", output_file_path=output_path) + copy_file_fixture(fixture_name="poster.jpg", output_file_path=output_thumbnail_path) return True elif "fanart" in output_name: - copy_file_fixture(fixture_name="fanart.jpeg", output_file_path=output_path) + copy_file_fixture(fixture_name="fanart.jpeg", output_file_path=output_thumbnail_path) return True return False - with patch.object( - YtDlpDownloader, - "_download_thumbnail", - new=lambda _, thumbnail_url, output_thumbnail_path: _mock_download_thumbnail( - output_thumbnail_path - ), + with patch( + "ytdl_sub.downloaders.downloader.download_and_convert_url_thumbnail", + new=_mock_download_and_convert_url_thumbnail, ): yield # TODO: create file here @@ -126,11 +126,9 @@ def mock_download_collection_entries( mock_download_collection_thumbnail, mock_entry_dict_factory: Callable, working_directory: str ): @contextlib.contextmanager - def _mock_download_collection_entries_factory(is_youtube_channel: bool): + def _mock_download_collection_entries_factory(is_youtube_channel: bool, num_urls: int = 1): def _write_entries_to_working_dir(*args, **kwargs) -> List[Dict]: - if (len(args[0].collection.urls.list) == 1) or ( - "season.2" in kwargs["url"] and len(args[0].download_options.urls.list) > 1 - ): + if num_urls == 1 or ("season.2" in kwargs["url"] and num_urls > 1): return [ mock_entry_dict_factory( uid="21-1", @@ -202,7 +200,7 @@ def mock_download_collection_entries( ] with patch.object( - YtDlpDownloader, "extract_info_via_info_json", new=_write_entries_to_working_dir + YTDLP, "extract_info_via_info_json", new=_write_entries_to_working_dir ), patch.object( YtDlpDownloader, "_extract_entry_info_with_retry", new=lambda _, entry: entry ): diff --git a/tests/unit/prebuilt_presets/test_prebuilt_presets.py b/tests/unit/prebuilt_presets/test_prebuilt_presets.py index 73c84020..4f761f95 100644 --- a/tests/unit/prebuilt_presets/test_prebuilt_presets.py +++ b/tests/unit/prebuilt_presets/test_prebuilt_presets.py @@ -215,7 +215,9 @@ class TestPrebuiltTvShowCollectionPresets: }, ) - with mock_download_collection_entries(is_youtube_channel=is_youtube_channel): + with mock_download_collection_entries( + is_youtube_channel=is_youtube_channel, num_urls=len(season_indices) + ): transaction_log = subscription.download(dry_run=False) assert_transaction_log_matches(