From 1922296c31465de76bb0274717f3d540f670a6cc Mon Sep 17 00:00:00 2001 From: Jesse Bannon Date: Mon, 13 Mar 2023 00:42:32 -0700 Subject: [PATCH] [REFACTOR] BaseDownloader class (#539) --- src/ytdl_sub/config/preset.py | 12 +++--- src/ytdl_sub/config/preset_class_mappings.py | 6 +-- src/ytdl_sub/downloaders/downloader.py | 43 +++++++++++++------ src/ytdl_sub/downloaders/generic/multi_url.py | 4 +- src/ytdl_sub/downloaders/generic/url.py | 4 +- .../subscriptions/base_subscription.py | 9 ++-- .../subscription_ytdl_options.py | 9 ++-- tests/unit/prebuilt_presets/conftest.py | 10 +++-- 8 files changed, 55 insertions(+), 42 deletions(-) diff --git a/src/ytdl_sub/config/preset.py b/src/ytdl_sub/config/preset.py index 8f47de81..5b365502 100644 --- a/src/ytdl_sub/config/preset.py +++ b/src/ytdl_sub/config/preset.py @@ -17,7 +17,7 @@ from ytdl_sub.config.preset_class_mappings import PluginMapping from ytdl_sub.config.preset_options import OutputOptions from ytdl_sub.config.preset_options import Overrides from ytdl_sub.config.preset_options import YTDLOptions -from ytdl_sub.downloaders.downloader import Downloader +from ytdl_sub.downloaders.downloader import BaseDownloader from ytdl_sub.downloaders.downloader import DownloaderValidator from ytdl_sub.entries.entry import Entry from ytdl_sub.plugins.plugin import Plugin @@ -123,7 +123,7 @@ class DownloadStrategyValidator(StrictDictValidator): validator=StringValidator, ).value - def get(self, downloader_source: str) -> Type[Downloader]: + def get(self, downloader_source: str) -> Type[BaseDownloader]: """ Parameters ---------- @@ -243,13 +243,13 @@ class Preset(_PresetShell): def _source_variables(self) -> List[str]: return Entry.source_variables() - def __validate_and_get_downloader(self, downloader_source: str) -> Type[Downloader]: + def __validate_and_get_downloader(self, downloader_source: str) -> Type[BaseDownloader]: return self._validate_key(key=downloader_source, validator=DownloadStrategyValidator).get( downloader_source=downloader_source ) def __validate_and_get_downloader_options( - self, downloader_source: str, downloader: Type[Downloader] + self, downloader_source: str, downloader: Type[BaseDownloader] ) -> DownloaderValidator: # Remove the download_strategy key before validating it against the downloader options # TODO: make this cleaner @@ -260,8 +260,8 @@ class Preset(_PresetShell): def __validate_and_get_downloader_and_options( self, - ) -> Tuple[Type[Downloader], DownloaderValidator]: - downloader: Optional[Type[Downloader]] = None + ) -> Tuple[Type[BaseDownloader], DownloaderValidator]: + downloader: Optional[Type[BaseDownloader]] = None download_options: Optional[DownloaderValidator] = None downloader_sources = DownloadStrategyMapping.sources() diff --git a/src/ytdl_sub/config/preset_class_mappings.py b/src/ytdl_sub/config/preset_class_mappings.py index b543a151..46d3439a 100644 --- a/src/ytdl_sub/config/preset_class_mappings.py +++ b/src/ytdl_sub/config/preset_class_mappings.py @@ -2,7 +2,7 @@ from typing import Dict from typing import List from typing import Type -from ytdl_sub.downloaders.downloader import Downloader +from ytdl_sub.downloaders.downloader import BaseDownloader from ytdl_sub.downloaders.generic.multi_url import MultiUrlDownloader from ytdl_sub.downloaders.generic.url import UrlDownloader from ytdl_sub.plugins.audio_extract import AudioExtractPlugin @@ -26,7 +26,7 @@ class DownloadStrategyMapping: Maps downloader strategies defined in the preset to its respective downloader class """ - _MAPPING: Dict[str, Dict[str, Type[Downloader]]] = { + _MAPPING: Dict[str, Dict[str, Type[BaseDownloader]]] = { "download": { "multi_url": MultiUrlDownloader, "url": UrlDownloader, @@ -81,7 +81,7 @@ class DownloadStrategyMapping: ) @classmethod - def get(cls, source: str, download_strategy: str) -> Type[Downloader]: + def get(cls, source: str, download_strategy: str) -> Type[BaseDownloader]: """ Parameters ---------- diff --git a/src/ytdl_sub/downloaders/downloader.py b/src/ytdl_sub/downloaders/downloader.py index 862d4927..194a52ab 100644 --- a/src/ytdl_sub/downloaders/downloader.py +++ b/src/ytdl_sub/downloaders/downloader.py @@ -100,18 +100,34 @@ class URLDownloadState: self.thumbnails_downloaded: Set[str] = set() -class Downloader(DownloadArchiver, Generic[DownloaderOptionsT], ABC): +class BaseDownloader(DownloadArchiver, Generic[DownloaderOptionsT], ABC): + downloader_options_type: Type[DownloaderValidator] = DownloaderValidator + + def __init__( + self, + download_options: DownloaderOptionsT, + enhanced_download_archive: EnhancedDownloadArchive, + download_ytdl_options: YTDLOptionsBuilder, + metadata_ytdl_options: YTDLOptionsBuilder, + overrides: Overrides, + ): + super().__init__(enhanced_download_archive=enhanced_download_archive) + self.download_options = download_options + self.overrides = overrides + self._download_ytdl_options_builder = download_ytdl_options + self._metadata_ytdl_options_builder = metadata_ytdl_options + + @abc.abstractmethod + def download(self) -> Iterable[Entry] | Iterable[Tuple[Entry, FileMetadata]]: + """The function to perform the download of all media entries""" + + +class YtDlpDownloader(BaseDownloader[DownloaderOptionsT], ABC): """ Class that interacts with ytdl to perform the download of metadata and content, and should translate that to list of Entry objects. """ - downloader_options_type: Type[DownloaderValidator] = DownloaderValidator - - supports_download_archive: bool = True - supports_subtitles: bool = True - supports_chapters: bool = True - _extract_entry_num_retries: int = 5 _extract_entry_retry_wait_sec: int = 5 @@ -147,13 +163,14 @@ class Downloader(DownloadArchiver, Generic[DownloaderOptionsT], ABC): overrides Override variables """ - DownloadArchiver.__init__(self=self, enhanced_download_archive=enhanced_download_archive) - self.download_options = download_options - self.overrides = overrides - self._download_ytdl_options_builder = download_ytdl_options - self._metadata_ytdl_options_builder = metadata_ytdl_options + super().__init__( + download_options=download_options, + enhanced_download_archive=enhanced_download_archive, + download_ytdl_options=download_ytdl_options, + metadata_ytdl_options=metadata_ytdl_options, + overrides=overrides, + ) self._downloaded_entries: Set[str] = set() - self._url_state: Optional[URLDownloadState] = None @property diff --git a/src/ytdl_sub/downloaders/generic/multi_url.py b/src/ytdl_sub/downloaders/generic/multi_url.py index a78d6f15..ce0ed70f 100644 --- a/src/ytdl_sub/downloaders/generic/multi_url.py +++ b/src/ytdl_sub/downloaders/generic/multi_url.py @@ -1,5 +1,5 @@ -from ytdl_sub.downloaders.downloader import Downloader from ytdl_sub.downloaders.downloader import DownloaderValidator +from ytdl_sub.downloaders.downloader import YtDlpDownloader from ytdl_sub.downloaders.generic.validators import MultiUrlValidator @@ -44,5 +44,5 @@ class MultiUrlDownloadOptions(MultiUrlValidator, DownloaderValidator): return self -class MultiUrlDownloader(Downloader[MultiUrlDownloadOptions]): +class MultiUrlDownloader(YtDlpDownloader[MultiUrlDownloadOptions]): downloader_options_type = MultiUrlDownloadOptions diff --git a/src/ytdl_sub/downloaders/generic/url.py b/src/ytdl_sub/downloaders/generic/url.py index bf29ef59..ed1091a2 100644 --- a/src/ytdl_sub/downloaders/generic/url.py +++ b/src/ytdl_sub/downloaders/generic/url.py @@ -1,5 +1,5 @@ -from ytdl_sub.downloaders.downloader import Downloader from ytdl_sub.downloaders.downloader import DownloaderValidator +from ytdl_sub.downloaders.downloader import YtDlpDownloader from ytdl_sub.downloaders.generic.validators import MultiUrlValidator from ytdl_sub.downloaders.generic.validators import UrlValidator @@ -36,5 +36,5 @@ class UrlDownloadOptions(UrlValidator, DownloaderValidator): ) -class UrlDownloader(Downloader[UrlDownloadOptions]): +class UrlDownloader(YtDlpDownloader[UrlDownloadOptions]): downloader_options_type = UrlDownloadOptions diff --git a/src/ytdl_sub/subscriptions/base_subscription.py b/src/ytdl_sub/subscriptions/base_subscription.py index 143b9c31..cf314d1d 100644 --- a/src/ytdl_sub/subscriptions/base_subscription.py +++ b/src/ytdl_sub/subscriptions/base_subscription.py @@ -8,7 +8,7 @@ from ytdl_sub.config.preset import PresetPlugins from ytdl_sub.config.preset_options import OutputOptions from ytdl_sub.config.preset_options import Overrides from ytdl_sub.config.preset_options import YTDLOptions -from ytdl_sub.downloaders.downloader import Downloader +from ytdl_sub.downloaders.downloader import BaseDownloader from ytdl_sub.downloaders.downloader import DownloaderValidator from ytdl_sub.ytdl_additions.enhanced_download_archive import EnhancedDownloadArchive @@ -52,7 +52,7 @@ class BaseSubscription(ABC): ) @property - def downloader_class(self) -> Type[Downloader]: + def downloader_class(self) -> Type[BaseDownloader]: """ Returns ------- @@ -130,10 +130,7 @@ class BaseSubscription(ABC): ------- Whether to maintain a download archive """ - return ( - self.output_options.maintain_download_archive - and self.downloader_class.supports_download_archive - ) + return self.output_options.maintain_download_archive @property def num_entries_added(self) -> int: diff --git a/src/ytdl_sub/subscriptions/subscription_ytdl_options.py b/src/ytdl_sub/subscriptions/subscription_ytdl_options.py index 420e4661..9acb2866 100644 --- a/src/ytdl_sub/subscriptions/subscription_ytdl_options.py +++ b/src/ytdl_sub/subscriptions/subscription_ytdl_options.py @@ -6,7 +6,7 @@ from typing import Type from typing import TypeVar from ytdl_sub.config.preset import Preset -from ytdl_sub.downloaders.downloader import Downloader +from ytdl_sub.downloaders.downloader import BaseDownloader from ytdl_sub.downloaders.ytdl_options_builder import YTDLOptionsBuilder from ytdl_sub.plugins.audio_extract import AudioExtractPlugin from ytdl_sub.plugins.chapters import ChaptersPlugin @@ -43,7 +43,7 @@ class SubscriptionYTDLOptions: return None @property - def _downloader(self) -> Type[Downloader]: + def _downloader(self) -> Type[BaseDownloader]: return self._preset.downloader @property @@ -82,10 +82,7 @@ class SubscriptionYTDLOptions: def _output_options(self) -> Dict: ytdl_options = {} - if ( - self._downloader.supports_download_archive - and self._preset.output_options.maintain_download_archive - ): + if self._preset.output_options.maintain_download_archive: ytdl_options["download_archive"] = str( Path(self._working_directory) / self._enhanced_download_archive.archive_file_name ) diff --git a/tests/unit/prebuilt_presets/conftest.py b/tests/unit/prebuilt_presets/conftest.py index e1f8ad14..a78af19f 100644 --- a/tests/unit/prebuilt_presets/conftest.py +++ b/tests/unit/prebuilt_presets/conftest.py @@ -11,7 +11,7 @@ import pytest from resources import copy_file_fixture from ytdl_sub.config.config_file import ConfigFile -from ytdl_sub.downloaders.downloader import Downloader +from ytdl_sub.downloaders.downloader import YtDlpDownloader from ytdl_sub.entries.variables.kwargs import DESCRIPTION from ytdl_sub.entries.variables.kwargs import EPOCH from ytdl_sub.entries.variables.kwargs import EXT @@ -112,7 +112,7 @@ def mock_download_collection_thumbnail(mock_downloaded_file_path): return False with patch.object( - Downloader, + YtDlpDownloader, "_download_thumbnail", new=lambda _, thumbnail_url, output_thumbnail_path: _mock_download_thumbnail( output_thumbnail_path @@ -202,8 +202,10 @@ def mock_download_collection_entries( ] with patch.object( - Downloader, "extract_info_via_info_json", new=_write_entries_to_working_dir - ), patch.object(Downloader, "_extract_entry_info_with_retry", new=lambda _, entry: entry): + YtDlpDownloader, "extract_info_via_info_json", new=_write_entries_to_working_dir + ), patch.object( + YtDlpDownloader, "_extract_entry_info_with_retry", new=lambda _, entry: entry + ): # Stub out metadata. TODO: update this if we do metadata plugins yield