[REFACTOR] BaseDownloader class (#539)

This commit is contained in:
Jesse Bannon 2023-03-13 00:42:32 -07:00 committed by GitHub
parent 52c1cc437d
commit 1922296c31
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
8 changed files with 55 additions and 42 deletions

View file

@ -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 OutputOptions
from ytdl_sub.config.preset_options import Overrides from ytdl_sub.config.preset_options import Overrides
from ytdl_sub.config.preset_options import YTDLOptions 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.downloaders.downloader import DownloaderValidator
from ytdl_sub.entries.entry import Entry from ytdl_sub.entries.entry import Entry
from ytdl_sub.plugins.plugin import Plugin from ytdl_sub.plugins.plugin import Plugin
@ -123,7 +123,7 @@ class DownloadStrategyValidator(StrictDictValidator):
validator=StringValidator, validator=StringValidator,
).value ).value
def get(self, downloader_source: str) -> Type[Downloader]: def get(self, downloader_source: str) -> Type[BaseDownloader]:
""" """
Parameters Parameters
---------- ----------
@ -243,13 +243,13 @@ class Preset(_PresetShell):
def _source_variables(self) -> List[str]: def _source_variables(self) -> List[str]:
return Entry.source_variables() 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( return self._validate_key(key=downloader_source, validator=DownloadStrategyValidator).get(
downloader_source=downloader_source downloader_source=downloader_source
) )
def __validate_and_get_downloader_options( def __validate_and_get_downloader_options(
self, downloader_source: str, downloader: Type[Downloader] self, downloader_source: str, downloader: Type[BaseDownloader]
) -> DownloaderValidator: ) -> DownloaderValidator:
# Remove the download_strategy key before validating it against the downloader options # Remove the download_strategy key before validating it against the downloader options
# TODO: make this cleaner # TODO: make this cleaner
@ -260,8 +260,8 @@ class Preset(_PresetShell):
def __validate_and_get_downloader_and_options( def __validate_and_get_downloader_and_options(
self, self,
) -> Tuple[Type[Downloader], DownloaderValidator]: ) -> Tuple[Type[BaseDownloader], DownloaderValidator]:
downloader: Optional[Type[Downloader]] = None downloader: Optional[Type[BaseDownloader]] = None
download_options: Optional[DownloaderValidator] = None download_options: Optional[DownloaderValidator] = None
downloader_sources = DownloadStrategyMapping.sources() downloader_sources = DownloadStrategyMapping.sources()

View file

@ -2,7 +2,7 @@ from typing import Dict
from typing import List from typing import List
from typing import Type 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.multi_url import MultiUrlDownloader
from ytdl_sub.downloaders.generic.url import UrlDownloader from ytdl_sub.downloaders.generic.url import UrlDownloader
from ytdl_sub.plugins.audio_extract import AudioExtractPlugin 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 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": { "download": {
"multi_url": MultiUrlDownloader, "multi_url": MultiUrlDownloader,
"url": UrlDownloader, "url": UrlDownloader,
@ -81,7 +81,7 @@ class DownloadStrategyMapping:
) )
@classmethod @classmethod
def get(cls, source: str, download_strategy: str) -> Type[Downloader]: def get(cls, source: str, download_strategy: str) -> Type[BaseDownloader]:
""" """
Parameters Parameters
---------- ----------

View file

@ -100,18 +100,34 @@ class URLDownloadState:
self.thumbnails_downloaded: Set[str] = set() 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, Class that interacts with ytdl to perform the download of metadata and content,
and should translate that to list of Entry objects. 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_num_retries: int = 5
_extract_entry_retry_wait_sec: int = 5 _extract_entry_retry_wait_sec: int = 5
@ -147,13 +163,14 @@ class Downloader(DownloadArchiver, Generic[DownloaderOptionsT], ABC):
overrides overrides
Override variables Override variables
""" """
DownloadArchiver.__init__(self=self, enhanced_download_archive=enhanced_download_archive) super().__init__(
self.download_options = download_options download_options=download_options,
self.overrides = overrides enhanced_download_archive=enhanced_download_archive,
self._download_ytdl_options_builder = download_ytdl_options download_ytdl_options=download_ytdl_options,
self._metadata_ytdl_options_builder = metadata_ytdl_options metadata_ytdl_options=metadata_ytdl_options,
overrides=overrides,
)
self._downloaded_entries: Set[str] = set() self._downloaded_entries: Set[str] = set()
self._url_state: Optional[URLDownloadState] = None self._url_state: Optional[URLDownloadState] = None
@property @property

View file

@ -1,5 +1,5 @@
from ytdl_sub.downloaders.downloader import Downloader
from ytdl_sub.downloaders.downloader import DownloaderValidator 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 MultiUrlValidator
@ -44,5 +44,5 @@ class MultiUrlDownloadOptions(MultiUrlValidator, DownloaderValidator):
return self return self
class MultiUrlDownloader(Downloader[MultiUrlDownloadOptions]): class MultiUrlDownloader(YtDlpDownloader[MultiUrlDownloadOptions]):
downloader_options_type = MultiUrlDownloadOptions downloader_options_type = MultiUrlDownloadOptions

View file

@ -1,5 +1,5 @@
from ytdl_sub.downloaders.downloader import Downloader
from ytdl_sub.downloaders.downloader import DownloaderValidator 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 MultiUrlValidator
from ytdl_sub.downloaders.generic.validators import UrlValidator 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 downloader_options_type = UrlDownloadOptions

View file

@ -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 OutputOptions
from ytdl_sub.config.preset_options import Overrides from ytdl_sub.config.preset_options import Overrides
from ytdl_sub.config.preset_options import YTDLOptions 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.downloaders.downloader import DownloaderValidator
from ytdl_sub.ytdl_additions.enhanced_download_archive import EnhancedDownloadArchive from ytdl_sub.ytdl_additions.enhanced_download_archive import EnhancedDownloadArchive
@ -52,7 +52,7 @@ class BaseSubscription(ABC):
) )
@property @property
def downloader_class(self) -> Type[Downloader]: def downloader_class(self) -> Type[BaseDownloader]:
""" """
Returns Returns
------- -------
@ -130,10 +130,7 @@ class BaseSubscription(ABC):
------- -------
Whether to maintain a download archive Whether to maintain a download archive
""" """
return ( return self.output_options.maintain_download_archive
self.output_options.maintain_download_archive
and self.downloader_class.supports_download_archive
)
@property @property
def num_entries_added(self) -> int: def num_entries_added(self) -> int:

View file

@ -6,7 +6,7 @@ from typing import Type
from typing import TypeVar from typing import TypeVar
from ytdl_sub.config.preset import Preset 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.downloaders.ytdl_options_builder import YTDLOptionsBuilder
from ytdl_sub.plugins.audio_extract import AudioExtractPlugin from ytdl_sub.plugins.audio_extract import AudioExtractPlugin
from ytdl_sub.plugins.chapters import ChaptersPlugin from ytdl_sub.plugins.chapters import ChaptersPlugin
@ -43,7 +43,7 @@ class SubscriptionYTDLOptions:
return None return None
@property @property
def _downloader(self) -> Type[Downloader]: def _downloader(self) -> Type[BaseDownloader]:
return self._preset.downloader return self._preset.downloader
@property @property
@ -82,10 +82,7 @@ class SubscriptionYTDLOptions:
def _output_options(self) -> Dict: def _output_options(self) -> Dict:
ytdl_options = {} ytdl_options = {}
if ( if self._preset.output_options.maintain_download_archive:
self._downloader.supports_download_archive
and self._preset.output_options.maintain_download_archive
):
ytdl_options["download_archive"] = str( ytdl_options["download_archive"] = str(
Path(self._working_directory) / self._enhanced_download_archive.archive_file_name Path(self._working_directory) / self._enhanced_download_archive.archive_file_name
) )

View file

@ -11,7 +11,7 @@ import pytest
from resources import copy_file_fixture from resources import copy_file_fixture
from ytdl_sub.config.config_file import ConfigFile 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 DESCRIPTION
from ytdl_sub.entries.variables.kwargs import EPOCH from ytdl_sub.entries.variables.kwargs import EPOCH
from ytdl_sub.entries.variables.kwargs import EXT from ytdl_sub.entries.variables.kwargs import EXT
@ -112,7 +112,7 @@ def mock_download_collection_thumbnail(mock_downloaded_file_path):
return False return False
with patch.object( with patch.object(
Downloader, YtDlpDownloader,
"_download_thumbnail", "_download_thumbnail",
new=lambda _, thumbnail_url, output_thumbnail_path: _mock_download_thumbnail( new=lambda _, thumbnail_url, output_thumbnail_path: _mock_download_thumbnail(
output_thumbnail_path output_thumbnail_path
@ -202,8 +202,10 @@ def mock_download_collection_entries(
] ]
with patch.object( with patch.object(
Downloader, "extract_info_via_info_json", new=_write_entries_to_working_dir YtDlpDownloader, "extract_info_via_info_json", new=_write_entries_to_working_dir
), patch.object(Downloader, "_extract_entry_info_with_retry", new=lambda _, entry: entry): ), patch.object(
YtDlpDownloader, "_extract_entry_info_with_retry", new=lambda _, entry: entry
):
# Stub out metadata. TODO: update this if we do metadata plugins # Stub out metadata. TODO: update this if we do metadata plugins
yield yield