ytdl-sub/src/ytdl_sub/downloaders/url/downloader.py
dependabot[bot] 64d3082a8a
[DEV] Bump pylint from 2.13.5 to 3.1.0 (#972)
* Bump pylint from 2.13.5 to 3.1.0

Bumps [pylint](https://github.com/pylint-dev/pylint) from 2.13.5 to 3.1.0.
- [Release notes](https://github.com/pylint-dev/pylint/releases)
- [Commits](https://github.com/pylint-dev/pylint/compare/v2.13.5...v3.1.0)

---
updated-dependencies:
- dependency-name: pylint
  dependency-type: direct:production
  update-type: version-update:semver-major
...

Signed-off-by: dependabot[bot] <support@github.com>

* bend knee to pylint

---------

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: Jesse Bannon <jbann1994@gmail.com>
2024-04-28 10:40:15 -07:00

543 lines
20 KiB
Python

import contextlib
import os
from pathlib import Path
from typing import Dict
from typing import Iterable
from typing import Iterator
from typing import List
from typing import Optional
from typing import Set
from typing import Tuple
from yt_dlp.utils import RejectedVideoReached
from ytdl_sub.config.overrides import Overrides
from ytdl_sub.downloaders.source_plugin import SourcePlugin
from ytdl_sub.downloaders.source_plugin import SourcePluginExtension
from ytdl_sub.downloaders.url.validators import MultiUrlValidator
from ytdl_sub.downloaders.url.validators import UrlThumbnailListValidator
from ytdl_sub.downloaders.url.validators import UrlValidator
from ytdl_sub.downloaders.ytdl_options_builder import YTDLOptionsBuilder
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.script.variable_definitions import VARIABLES
from ytdl_sub.entries.script.variable_definitions import VariableDefinitions
from ytdl_sub.utils.file_handler import FileHandler
from ytdl_sub.utils.logger import Logger
from ytdl_sub.utils.script import ScriptUtils
from ytdl_sub.utils.thumbnail import ThumbnailTypes
from ytdl_sub.utils.thumbnail import download_and_convert_url_thumbnail
from ytdl_sub.utils.thumbnail import try_convert_download_thumbnail
from ytdl_sub.ytdl_additions.enhanced_download_archive import EnhancedDownloadArchive
v: VariableDefinitions = VARIABLES
download_logger = Logger.get(name="downloader")
class URLDownloadState:
def __init__(self, entries_total: int):
self.entries_total = entries_total
self.entries_downloaded = 0
class UrlDownloaderBasePluginExtension(SourcePluginExtension[MultiUrlValidator]):
def _match_entry_to_url_validator(self, entry: Entry) -> UrlValidator:
"""
Handle matching a URL to its original validator. This is for .info.json updates
when older entries have missing variables
"""
input_url_idx = entry.get(v.ytdl_sub_input_url_index, int)
entry_input_url = entry.get(v.ytdl_sub_input_url, str)
if 0 <= input_url_idx < len(self.plugin_options.urls.list):
validator = self.plugin_options.urls.list[input_url_idx]
if self.overrides.apply_formatter(validator.url) == entry_input_url:
return validator
# Match the first validator based on the URL, if one exists
for validator in self.plugin_options.urls.list:
if self.overrides.apply_formatter(validator.url) == entry_input_url:
return validator
# Return the first validator if none exist
return self.plugin_options.urls.list[0]
class UrlDownloaderThumbnailPlugin(UrlDownloaderBasePluginExtension):
def __init__(
self,
options: MultiUrlValidator,
overrides: Overrides,
enhanced_download_archive: EnhancedDownloadArchive,
):
super().__init__(
options=options,
overrides=overrides,
enhanced_download_archive=enhanced_download_archive,
)
self._thumbnails_downloaded: Set[str] = set()
def _download_parent_thumbnails(
self,
thumbnail_list_info: UrlThumbnailListValidator,
entry: Entry,
parent: EntryParent,
) -> None:
"""
Downloads and moves channel avatar and banner images to the output directory.
"""
for thumbnail_info in thumbnail_list_info.list:
thumbnail_name = self.overrides.apply_formatter(thumbnail_info.name, entry=entry)
thumbnail_id = self.overrides.apply_formatter(thumbnail_info.uid)
# If the thumbnail name is an empty string, completely ignore trying to download it
if not thumbnail_name:
continue
# If latest entry, always update the thumbnail on each entry
if thumbnail_id == ThumbnailTypes.LATEST_ENTRY:
# always save in dry-run even if it doesn't exist...
if self.is_dry_run or entry.is_thumbnail_downloaded():
self.save_file(
file_name=entry.get_download_thumbnail_name(),
output_file_name=thumbnail_name,
copy_file=True,
)
self._thumbnails_downloaded.add(thumbnail_name)
continue
# If not latest entry and the thumbnail has already been downloaded, then skip
if thumbnail_name in self._thumbnails_downloaded:
continue
if (thumbnail_url := parent.get_thumbnail_url(thumbnail_id=thumbnail_id)) is None:
download_logger.debug("Failed to find thumbnail id '%s'", thumbnail_id)
continue
if download_and_convert_url_thumbnail(
thumbnail_url=thumbnail_url,
output_thumbnail_path=str(Path(self.working_directory) / thumbnail_name),
):
self.save_file(file_name=thumbnail_name)
self._thumbnails_downloaded.add(thumbnail_name)
else:
download_logger.debug("Failed to download thumbnail id '%s'", thumbnail_id)
def _download_url_thumbnails(self, collection_url: UrlValidator, entry: Entry):
"""
After all media entries have been downloaded, post processed, and moved to the output
directory, run this function. This lets the downloader add any extra files directly to the
output directory, for things like YT channel image, banner.
"""
if source_metadata := entry.get(v.source_metadata, dict):
self._download_parent_thumbnails(
thumbnail_list_info=collection_url.source_thumbnails,
entry=entry,
parent=EntryParent(source_metadata, working_directory=self.working_directory),
)
if playlist_metadata := entry.get(v.playlist_metadata, dict):
self._download_parent_thumbnails(
thumbnail_list_info=collection_url.playlist_thumbnails,
entry=entry,
parent=EntryParent(playlist_metadata, working_directory=self.working_directory),
)
def modify_entry(self, entry: Entry) -> Optional[Entry]:
"""
Use the entry to download thumbnails (or move if LATEST_ENTRY).
In addition, convert the entry thumbnail to jpg
"""
# We always convert entry thumbnails to jpgs, and is performed here to be done
# as early as possible in the plugin pipeline (downstream plugins depend on it being jpg)
if not self.is_dry_run:
try_convert_download_thumbnail(entry=entry)
self._download_url_thumbnails(
collection_url=self._match_entry_to_url_validator(entry=entry),
entry=entry,
)
return entry
class UrlDownloaderCollectionVariablePlugin(UrlDownloaderBasePluginExtension):
def __init__(
self,
options: MultiUrlValidator,
overrides: Overrides,
enhanced_download_archive: EnhancedDownloadArchive,
):
super().__init__(
options=options,
overrides=overrides,
enhanced_download_archive=enhanced_download_archive,
)
self._thumbnails_downloaded: Set[str] = set()
def modify_entry_metadata(self, entry: Entry) -> Optional[Entry]:
"""
Add collection variables to the entry
"""
collection_url = self._match_entry_to_url_validator(entry=entry)
entry.add(collection_url.variables.dict_with_format_strings)
return entry
class MultiUrlDownloader(SourcePlugin[MultiUrlValidator]):
"""
Class that interacts with ytdl to perform the download of metadata and content,
and should translate that to list of Entry objects.
"""
plugin_options_type = MultiUrlValidator
plugin_extensions = [UrlDownloaderThumbnailPlugin, UrlDownloaderCollectionVariablePlugin]
@classmethod
def ytdl_option_defaults(cls) -> Dict:
"""
.. code-block:: yaml
ytdl_options:
ignoreerrors: True # ignore errors like hidden videos, age restriction, etc
"""
return {"ignoreerrors": True}
def __init__(
self,
options: MultiUrlValidator,
enhanced_download_archive: EnhancedDownloadArchive,
download_ytdl_options: YTDLOptionsBuilder,
metadata_ytdl_options: YTDLOptionsBuilder,
overrides: Overrides,
):
"""
Parameters
----------
options
Options validator for this downloader
enhanced_download_archive
Download archive
download_ytdl_options
YTDL options builder for downloading media
metadata_ytdl_options
YTDL options builder for downloading metadata
overrides
Override variables
"""
super().__init__(
options=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
def download_ytdl_options(self, url_idx: Optional[int] = None) -> Dict:
"""
Returns
-------
YTLD options dict for downloading
"""
return (
self._download_ytdl_options_builder.clone()
.add(self.ytdl_option_defaults(), before=True)
.add(
(
self.plugin_options.urls.list[url_idx].ytdl_options.dict
if url_idx is not None
else None
),
before=True,
)
.to_dict()
)
def metadata_ytdl_options(self, ytdl_option_overrides: Dict) -> Dict:
"""
Returns
-------
YTDL options dict for fetching metadata
"""
return (
self._metadata_ytdl_options_builder.clone()
.add(self.ytdl_option_defaults(), before=True)
.add(ytdl_option_overrides, before=True)
.to_dict()
)
@property
def is_dry_run(self) -> bool:
"""
Returns
-------
True if dry-run is enabled. False otherwise.
"""
return self.download_ytdl_options().get("skip_download", False)
@property
def is_entry_thumbnails_enabled(self) -> bool:
"""
Returns
-------
True if entry thumbnails should be downloaded. False otherwise.
"""
return self.download_ytdl_options().get("writethumbnail", False)
###############################################################################################
# DOWNLOAD FUNCTIONS
def _is_downloaded(self, entry: Entry) -> bool:
return entry.ytdl_uid() in self._downloaded_entries
def _mark_downloaded(self, entry: Entry) -> None:
self._downloaded_entries.add(entry.ytdl_uid())
@property
def collection(self) -> MultiUrlValidator:
"""Return the download options collection"""
return self.plugin_options
@contextlib.contextmanager
def _separate_download_archives(self, clear_info_json_files: bool = False):
"""
Separate download archive writing between collection urls. This is so break_on_existing
does not break when downloading from subset urls.
Parameters
----------
clear_info_json_files
Whether to delete info.json files after yield
"""
archive_path = self.download_ytdl_options().get("download_archive", "")
backup_archive_path = f"{archive_path}.backup"
# If archive path exists, maintain download archive is enable
if archive_file_exists := archive_path and os.path.isfile(archive_path):
archive_file_exists = True
# If a backup exists, it's the one prior to any downloading, use that.
if os.path.isfile(backup_archive_path):
FileHandler.copy(src_file_path=backup_archive_path, dst_file_path=archive_path)
# If not, create the backup
else:
FileHandler.copy(src_file_path=archive_path, dst_file_path=backup_archive_path)
yield
# If an archive path did not exist at first, but now exists, delete it
if not archive_file_exists and os.path.isfile(archive_path):
FileHandler.delete(file_path=archive_path)
# If the archive file did exist, restore the backup
elif archive_file_exists:
FileHandler.move(src_file_path=backup_archive_path, dst_file_path=archive_path)
# Clear info json files if true
if clear_info_json_files:
info_json_files = [
Path(self.working_directory) / path
for path in os.listdir(self.working_directory)
if path.endswith(".info.json")
]
for info_json_file in info_json_files:
FileHandler.delete(info_json_file)
def _extract_entry_info_with_retry(self, entry: Entry) -> Entry:
download_entry_dict = YTDLP.extract_info_with_retry(
ytdl_options_overrides=self.download_ytdl_options(
url_idx=entry.get(v.ytdl_sub_input_url_index, int)
),
is_downloaded_fn=None if self.is_dry_run else entry.is_downloaded,
is_thumbnail_downloaded_fn=(
None
if (self.is_dry_run or not self.is_entry_thumbnails_enabled)
else entry.is_thumbnail_downloaded_via_ytdlp
),
url=entry.webpage_url,
)
return Entry(
download_entry_dict,
working_directory=self.working_directory,
)
def _iterate_child_entries(
self, entries: List[Entry], download_reversed: bool
) -> Iterator[Entry]:
# Iterate a list of entries, and delete the entries after yielding
entries_to_iter: List[Optional[Entry]] = entries
indices = list(range(len(entries_to_iter)))
if download_reversed:
indices = reversed(indices)
for idx in indices:
self._url_state.entries_downloaded += 1
if self._is_downloaded(entries_to_iter[idx]):
download_logger.info(
"Already downloaded entry %d/%d: %s",
self._url_state.entries_downloaded,
self._url_state.entries_total,
entries_to_iter[idx].title,
)
entries_to_iter[idx] = None
continue
yield entries_to_iter[idx]
self._mark_downloaded(entries_to_iter[idx])
entries_to_iter[idx] = None
def _iterate_parent_entry(
self, parent: EntryParent, download_reversed: bool
) -> Iterator[Entry]:
yield from self._iterate_child_entries(
entries=parent.entry_children(), download_reversed=download_reversed
)
# Recursion the parent's parent entries
for parent_child in reversed(parent.parent_children()):
yield from self._iterate_parent_entry(
parent=parent_child, download_reversed=download_reversed
)
def _download_url_metadata(
self, url: str, include_sibling_metadata: bool, ytdl_options_overrides: Dict
) -> Tuple[List[EntryParent], List[Entry]]:
"""
Downloads only info.json files and forms EntryParent trees
"""
with self._separate_download_archives():
entry_dicts = YTDLP.extract_info_via_info_json(
working_directory=self.working_directory,
ytdl_options_overrides=ytdl_options_overrides,
log_prefix_on_info_json_dl="Downloading metadata for",
url=url,
)
parents = EntryParent.from_entry_dicts(
url=url,
entry_dicts=entry_dicts,
working_directory=self.working_directory,
include_sibling_metadata=include_sibling_metadata,
)
orphans = EntryParent.from_entry_dicts_with_no_parents(
parents=parents,
entry_dicts=entry_dicts,
working_directory=self.working_directory,
)
return parents, orphans
def _iterate_entries(
self,
parents: List[EntryParent],
orphans: List[Entry],
download_reversed: bool,
) -> Iterator[Entry]:
"""
Downloads the leaf entries from EntryParent trees
"""
# Delete info json files afterwards so other collection URLs do not use them
with self._separate_download_archives(clear_info_json_files=True):
for parent in parents:
yield from self._iterate_parent_entry(
parent=parent, download_reversed=download_reversed
)
yield from self._iterate_child_entries(
entries=orphans, download_reversed=download_reversed
)
def _download_metadata(self, url: str, validator: UrlValidator) -> Iterable[Entry]:
metadata_ytdl_options = self.metadata_ytdl_options(
ytdl_option_overrides=validator.ytdl_options.dict
)
download_reversed = ScriptUtils.bool_formatter_output(
self.overrides.apply_formatter(validator.download_reverse)
)
parents, orphan_entries = self._download_url_metadata(
url=url,
include_sibling_metadata=validator.include_sibling_metadata,
ytdl_options_overrides=metadata_ytdl_options,
)
# TODO: Encapsulate this logic into its own class
self._url_state = URLDownloadState(
entries_total=sum(parent.num_children() for parent in parents) + len(orphan_entries)
)
download_logger.info("Beginning downloads for %s", url)
yield from self._iterate_entries(
parents=parents,
orphans=orphan_entries,
download_reversed=download_reversed,
)
def download_metadata(self) -> Iterable[Entry]:
"""The function to perform the download of all media entries"""
# download the bottom-most urls first since they are top-priority
for idx, url_validator in reversed(list(enumerate(self.collection.urls.list))):
# URLs can be empty. If they are, then skip
if not (url := self.overrides.apply_formatter(url_validator.url)):
continue
for entry in self._download_metadata(url=url, validator=url_validator):
entry.initialize_script(self.overrides).add(
{
v.ytdl_sub_input_url: url,
v.ytdl_sub_input_url_index: idx,
v.ytdl_sub_input_url_count: len(self.collection.urls.list),
}
)
yield entry
def download(self, entry: Entry) -> Optional[Entry]:
"""
Parameters
----------
entry
Entry to download
Returns
-------
The entry that was downloaded successfully
Raises
------
RejectedVideoReached
If a video was rejected and was not from match_filter
"""
download_logger.info(
"Downloading entry %d/%d: %s",
self._url_state.entries_downloaded,
self._url_state.entries_total,
entry.title,
)
# Match-filters can be applied at the download stage. If the download is rejected,
# then return None
try:
download_entry = self._extract_entry_info_with_retry(entry=entry)
except RejectedVideoReached:
download_logger.info("Entry rejected by download match-filter, skipping ..")
return None
upload_date_idx = self._enhanced_download_archive.mapping.get_num_entries_with_upload_date(
upload_date_standardized=entry.get(v.upload_date_standardized, str)
)
download_idx = self._enhanced_download_archive.num_entries
return entry.add_injected_variables(
download_entry=download_entry,
download_idx=download_idx,
upload_date_idx=upload_date_idx,
)