From c66e5b72f3bfdff36c0c2fc5045154b1921af0d0 Mon Sep 17 00:00:00 2001 From: Jesse Bannon Date: Fri, 16 Sep 2022 22:56:56 -0700 Subject: [PATCH] [BACKEND] Add source and playlist variables to every entry (#236) --- src/ytdl_sub/config/preset_class_mappings.py | 4 + src/ytdl_sub/downloaders/downloader.py | 75 ++++--- src/ytdl_sub/downloaders/youtube/channel.py | 9 +- src/ytdl_sub/downloaders/youtube/playlist.py | 11 +- src/ytdl_sub/downloaders/youtube/video.py | 4 +- src/ytdl_sub/entries/base_entry.py | 96 +++++++- src/ytdl_sub/entries/entry.py | 16 -- src/ytdl_sub/entries/entry_parent.py | 224 +++++++++++++++---- 8 files changed, 333 insertions(+), 106 deletions(-) diff --git a/src/ytdl_sub/config/preset_class_mappings.py b/src/ytdl_sub/config/preset_class_mappings.py index a129bf01..73be4059 100644 --- a/src/ytdl_sub/config/preset_class_mappings.py +++ b/src/ytdl_sub/config/preset_class_mappings.py @@ -3,6 +3,7 @@ from typing import List from typing import Type from ytdl_sub.downloaders.downloader import Downloader +from ytdl_sub.downloaders.generic.collection import CollectionDownloader from ytdl_sub.downloaders.soundcloud.albums_and_singles import SoundcloudAlbumsAndSinglesDownloader from ytdl_sub.downloaders.youtube.channel import YoutubeChannelDownloader from ytdl_sub.downloaders.youtube.merge_playlist import YoutubeMergePlaylistDownloader @@ -36,6 +37,9 @@ class DownloadStrategyMapping: "soundcloud": { "albums_and_singles": SoundcloudAlbumsAndSinglesDownloader, }, + "generic": { + "collection": CollectionDownloader, + }, } @classmethod diff --git a/src/ytdl_sub/downloaders/downloader.py b/src/ytdl_sub/downloaders/downloader.py index 0f466110..9ed979bf 100644 --- a/src/ytdl_sub/downloaders/downloader.py +++ b/src/ytdl_sub/downloaders/downloader.py @@ -47,20 +47,6 @@ def _entry_key(entry: BaseEntry) -> str: return entry.extractor + entry.uid -def _get_parent_entry_variables(parent: EntryParent) -> Dict[str, str | int]: - """ - Adds source variables to the child entry derived from the parent entry. - """ - if not parent.child_entries: - return {} - - return { - "playlist_max_upload_year": max( - child_entry.to_type(Entry).upload_year for child_entry in parent.child_entries - ) - } - - class DownloaderValidator(StrictDictValidator, AddsVariablesMixin, ABC): """ Placeholder class to define downloader options @@ -404,23 +390,34 @@ class Downloader(DownloadArchiver, Generic[DownloaderOptionsT, DownloaderEntryT] return entry def _download_parent_entry(self, parent: EntryParent) -> Generator[Entry, None, None]: - """Download in reverse order, that way we download older entries ones first""" - if parent.is_entry(): - yield self._download_entry(parent.to_type(Entry)) - return - + # Download the parent's entries first, in reverse order for entry_child in reversed(parent.entry_children()): if _entry_key(entry_child) in self.downloaded_entries: continue - yield self._download_entry(entry_child.to_type(Entry)) + yield self._download_entry(entry_child) self.downloaded_entries.add(_entry_key(entry_child)) + # Recursion the parent's parent entries for parent_child in reversed(parent.parent_children()): for entry_child in self._download_parent_entry(parent=parent_child): yield entry_child - def _download_url_metadata(self, collection_url: CollectionUrlValidator) -> List[EntryParent]: + def _set_collection_variables( + self, collection_url: CollectionUrlValidator, entry: Entry | EntryParent + ): + if isinstance(entry, EntryParent): + for child in entry.parent_children(): + self._set_collection_variables(collection_url, child) + for child in entry.entry_children(): + child.add_variables(variables_to_add=collection_url.variables) + + elif isinstance(entry, Entry): + entry.add_variables(variables_to_add=collection_url.variables) + + def _download_url_metadata( + self, collection_url: CollectionUrlValidator + ) -> Tuple[List[EntryParent], List[Entry]]: """ Downloads only info.json files and forms EntryParent trees """ @@ -432,32 +429,50 @@ class Downloader(DownloadArchiver, Generic[DownloaderOptionsT, DownloaderEntryT] ) self.parents = EntryParent.from_entry_dicts( - entry_dicts=entry_dicts, working_directory=self.working_directory + url=collection_url.url, + entry_dicts=entry_dicts, + working_directory=self.working_directory, + ) + orphans = EntryParent.from_entry_dicts_with_no_parents( + parents=self.parents, entry_dicts=entry_dicts, working_directory=self.working_directory ) - return self.parents - def _download_url( - self, collection_url: CollectionUrlValidator, parents: List[EntryParent] + for parent_entry in self.parents: + self._set_collection_variables(collection_url, parent_entry) + for entry in orphans: + self._set_collection_variables(collection_url, entry) + + return self.parents, orphans + + def _download( + self, + parents: Optional[List[EntryParent]] = None, + orphans: Optional[List[Entry]] = None, ) -> Generator[Entry, None, None]: """ Downloads the leaf entries from EntryParent trees """ + if parents is None: + parents = [] + if orphans is None: + orphans = [] + with self._separate_download_archives(): for parent in parents: for entry_child in self._download_parent_entry(parent=parent): - entry_child.add_variables( - dict(_get_parent_entry_variables(parent), **collection_url.variables) - ) yield entry_child + for orphan in orphans: + yield self._download_entry(orphan) + def download( self, ) -> Iterable[DownloaderEntryT] | Iterable[Tuple[DownloaderEntryT, FileMetadata]]: """The function to perform the download of all media entries""" # download the bottom-most urls first since they are top-priority for collection_url in reversed(self.collection.collection_urls.list): - parents = self._download_url_metadata(collection_url=collection_url) - for entry in self._download_url(collection_url=collection_url, parents=parents): + parents, orphan_entries = self._download_url_metadata(collection_url=collection_url) + for entry in self._download(parents=parents, orphans=orphan_entries): yield entry def post_download(self): diff --git a/src/ytdl_sub/downloaders/youtube/channel.py b/src/ytdl_sub/downloaders/youtube/channel.py index de684943..2db3a6cf 100644 --- a/src/ytdl_sub/downloaders/youtube/channel.py +++ b/src/ytdl_sub/downloaders/youtube/channel.py @@ -136,7 +136,8 @@ class YoutubeChannelDownloader(YoutubeDownloader[YoutubeChannelDownloaderOptions Downloads all videos from a channel """ collection_url = self.collection.collection_urls.list[0] - super()._download_url_metadata(collection_url=collection_url) + _, orphans = super()._download_url_metadata(collection_url=collection_url) + assert not orphans # TODO: Handle this better self.overrides.add_override_variables( @@ -147,10 +148,8 @@ class YoutubeChannelDownloader(YoutubeDownloader[YoutubeChannelDownloaderOptions } ) - for entry in super()._download_url(collection_url=collection_url, parents=self.parents): - # pylint: disable=protected-access - yield YoutubeVideo(entry_dict=entry._kwargs, working_directory=self.working_directory) - # pylint: enable=protected-access + for entry in super()._download(parents=self.parents): + yield entry.to_type(YoutubeVideo) def _download_thumbnail( self, diff --git a/src/ytdl_sub/downloaders/youtube/playlist.py b/src/ytdl_sub/downloaders/youtube/playlist.py index 68af8d2b..a3aed608 100644 --- a/src/ytdl_sub/downloaders/youtube/playlist.py +++ b/src/ytdl_sub/downloaders/youtube/playlist.py @@ -101,7 +101,8 @@ class YoutubePlaylistDownloader( Downloads all videos in a Youtube playlist. """ collection_url = self.collection.collection_urls.list[0] - super()._download_url_metadata(collection_url) + _, orphans = super()._download_url_metadata(collection_url) + assert not orphans # TODO: Handle this better self.overrides.add_override_variables( @@ -112,9 +113,5 @@ class YoutubePlaylistDownloader( } ) - for entry in super()._download_url(collection_url=collection_url, parents=self.parents): - # pylint: disable=protected-access - yield YoutubePlaylistVideo( - entry_dict=entry._kwargs, working_directory=self.working_directory - ) - # pylint: enable=protected-access + for entry in super()._download(parents=self.parents): + yield entry.to_type(YoutubePlaylistVideo) diff --git a/src/ytdl_sub/downloaders/youtube/video.py b/src/ytdl_sub/downloaders/youtube/video.py index a7813de2..3bee57eb 100644 --- a/src/ytdl_sub/downloaders/youtube/video.py +++ b/src/ytdl_sub/downloaders/youtube/video.py @@ -75,6 +75,4 @@ class YoutubeVideoDownloader(YoutubeDownloader[YoutubeVideoDownloaderOptions, Yo def download(self) -> List[YoutubeVideo]: """Downloads the single video""" for entry in super().download(): - # pylint: disable=protected-access - yield YoutubeVideo(entry_dict=entry._kwargs, working_directory=self.working_directory) - # pylint: enable=protected-access + yield entry.to_type(YoutubeVideo) diff --git a/src/ytdl_sub/entries/base_entry.py b/src/ytdl_sub/entries/base_entry.py index e86cf2f6..4c5e8248 100644 --- a/src/ytdl_sub/entries/base_entry.py +++ b/src/ytdl_sub/entries/base_entry.py @@ -1,8 +1,11 @@ from abc import ABC +from pathlib import Path from typing import Any from typing import Dict from typing import List from typing import Optional +from typing import Type +from typing import TypeVar from typing import final from yt_dlp.utils import sanitize_filename @@ -45,9 +48,9 @@ class BaseEntryVariables: Returns ------- str - The title of the entry + The title of the entry. If a title does not exist, returns its unique ID. """ - return self.kwargs("title") + return self.kwargs_get("title", self.uid) @property def title_sanitized(self) -> str: @@ -83,6 +86,9 @@ class BaseEntryVariables: # pylint: enable=no-member +TBaseEntry = TypeVar("TBaseEntry", bound="BaseEntry") + + class BaseEntry(BaseEntryVariables, ABC): """ Abstract entry object to represent anything download from ytdl (playlist metadata, media, etc). @@ -104,6 +110,20 @@ class BaseEntry(BaseEntryVariables, ABC): self._additional_variables: Dict[str, str | int] = {} + def base_variable_dict(self) -> Dict[str, str]: + """ + Returns + ------- + BaseEntry variables that can be nested for playlist, source, etc + """ + return { + "uid": self.uid, + "extractor": self.extractor, + "title": self.title, + "title_sanitized": self.title_sanitized, + "webpage_url": self.webpage_url, + } + def kwargs_contains(self, key: str) -> bool: """Returns whether internal kwargs contains the specified key""" return key in self._kwargs @@ -135,6 +155,23 @@ class BaseEntry(BaseEntryVariables, ABC): """ return self._working_directory + def add_kwargs(self, variables_to_add: Dict[str, Any]) -> "BaseEntry": + """ + Adds variables to kwargs. Use with caution since yt-dlp data can be overwritten. + Plugins should use ``add_variables``. + + Parameters + ---------- + variables_to_add + Variables to add to kwargs + + Returns + ------- + self + """ + self._kwargs = dict(self._kwargs, **variables_to_add) + return self + def add_variables(self, variables_to_add: Dict[str, str]) -> "BaseEntry": """ Parameters @@ -160,6 +197,22 @@ class BaseEntry(BaseEntryVariables, ABC): self._additional_variables = dict(self._additional_variables, **variables_to_add) return self + def get_download_info_json_name(self) -> str: + """ + Returns + ------- + The download info json's file name + """ + return f"{self.uid}.{self.info_json_ext}" + + def get_download_info_json_path(self) -> str: + """ + Returns + ------- + Entry's downloaded info json file path + """ + return str(Path(self.working_directory()) / self.get_download_info_json_name()) + def _added_variables(self) -> Dict[str, str]: """ Returns @@ -189,3 +242,42 @@ class BaseEntry(BaseEntryVariables, ABC): source_var: getattr(self, source_var) for source_var in self.source_variables() } return dict(source_variable_dict, **self._added_variables()) + + @final + def to_type(self, entry_type: Type[TBaseEntry]) -> TBaseEntry: + """ + Returns + ------- + Converted EntryParent to Entry-like class + """ + return entry_type(entry_dict=self._kwargs, working_directory=self._working_directory) + + @classmethod + def is_entry_parent(cls, entry_dict: Dict | TBaseEntry): + """ + Returns + ------- + True if it is a parent. False otherwise + """ + entry_type: Optional[str] = None + if isinstance(entry_dict, cls): + entry_type = entry_dict.kwargs_get("_type") + if isinstance(entry_dict, dict): + entry_type = entry_dict.get("_type") + + return entry_type == "playlist" + + @classmethod + def is_entry(cls, entry_dict: Dict | TBaseEntry): + """ + Returns + ------- + True if it is an entry. False otherwise + """ + entry_ext: Optional[str] = None + if isinstance(entry_dict, cls): + entry_ext = entry_dict.kwargs_get("ext") + if isinstance(entry_dict, dict): + entry_ext = entry_dict.get("ext") + + return entry_ext is not None diff --git a/src/ytdl_sub/entries/entry.py b/src/ytdl_sub/entries/entry.py index dfa130ce..12d44367 100644 --- a/src/ytdl_sub/entries/entry.py +++ b/src/ytdl_sub/entries/entry.py @@ -57,22 +57,6 @@ class Entry(EntryVariables, BaseEntry): return None - def get_download_info_json_name(self) -> str: - """ - Returns - ------- - The download info json's file name - """ - return f"{self.uid}.{self.info_json_ext}" - - def get_download_info_json_path(self) -> str: - """ - Returns - ------- - Entry's downloaded info json file path - """ - return str(Path(self.working_directory()) / self.get_download_info_json_name()) - def write_info_json(self) -> None: """ Write the entry's _kwargs back into the info.json file as well as its source variables diff --git a/src/ytdl_sub/entries/entry_parent.py b/src/ytdl_sub/entries/entry_parent.py index a9adebc9..06e7c674 100644 --- a/src/ytdl_sub/entries/entry_parent.py +++ b/src/ytdl_sub/entries/entry_parent.py @@ -1,12 +1,24 @@ +import functools +import math from typing import Dict from typing import List from typing import Optional -from typing import Type -from typing import TypeVar + +import mergedeep from ytdl_sub.entries.base_entry import BaseEntry +from ytdl_sub.entries.base_entry import TBaseEntry +from ytdl_sub.entries.entry import Entry -TBaseEntry = TypeVar("TBaseEntry", bound=BaseEntry) + +class ParentType: + PLAYLIST = "playlist" + SOURCE = "source" + + +def _sort_entries(entries: List[TBaseEntry]) -> List[TBaseEntry]: + """Try sorting by playlist_id first, then fall back to uid""" + return sorted(entries, key=lambda ent: (ent.kwargs_get("playlist_index", math.inf), ent.uid)) class EntryParent(BaseEntry): @@ -14,49 +26,106 @@ class EntryParent(BaseEntry): super().__init__(entry_dict=entry_dict, working_directory=working_directory) self.child_entries: List["EntryParent"] = [] - def is_entry(self) -> bool: - """ - Returns - ------- - True if the entry contains a media file. False otherwise. - """ - return self.kwargs_contains("ext") - + @functools.cache def parent_children(self) -> List["EntryParent"]: """This parent's children that are also parents""" - return [child for child in self.child_entries if child.child_count() > 0] + return _sort_entries([child for child in self.child_entries if self.is_entry_parent(child)]) - def entry_children(self) -> List["EntryParent"]: + @functools.cache + def entry_children(self) -> List[Entry]: """This parent's children that are entries""" - return [child for child in self.child_entries if child.is_entry()] + return _sort_entries( + [child.to_type(Entry) for child in self.child_entries if self.is_entry(child)] + ) - def read_children_from_entry_dicts(self, entry_dicts: List[Dict]) -> "EntryParent": + def _playlist_variables(self, idx: int, children: List[TBaseEntry], parent_type: str) -> Dict: + _count = self.kwargs_get("playlist_count", len(children)) + _index = children[idx].kwargs_get("playlist_index", idx + 1) + + if parent_type == ParentType.SOURCE: + return {"source_index": _index, "source_count": _count} + return { + "source_index": self.kwargs_get("source_index", 1), + "source_count": self.kwargs_get("source_count", 1), + "playlist_index": _index, + "playlist_count": _count, + } + + def _parent_variables(self, parent_type: str) -> Dict: + return dict( + {f"{parent_type}_entry": self._kwargs}, + **{f"{parent_type}_{key}": value for key, value in self.base_variable_dict().items()}, + ) + + def _get_entry_children_variable_list(self, variable_name: str) -> List[str | int]: + return [getattr(entry_child, variable_name) for entry_child in self.entry_children()] + + def _entry_aggregate_variables(self) -> Dict: + if not self.entry_children(): + return {} + + return { + "playlist_max_upload_year": max(self._get_entry_children_variable_list("upload_year")), + "playlist_max_upload_year_truncated": max( + self._get_entry_children_variable_list("upload_year_truncated") + ), + } + + # pylint: disable=protected-access + + def _set_child_variables(self, parents: Optional[List["EntryParent"]] = None) -> "EntryParent": + if parents is None: + parents = [self] + self.add_kwargs( + self._playlist_variables(idx=0, children=parents, parent_type=ParentType.SOURCE) + ) + + kwargs_to_add: Dict = {} + if len(parents) >= 1: + mergedeep.merge(kwargs_to_add, parents[-1]._parent_variables(ParentType.PLAYLIST)) + if len(parents) >= 2: + mergedeep.merge(kwargs_to_add, parents[-2]._parent_variables(ParentType.SOURCE)) + if len(parents) >= 3: + raise ValueError( + "ytdl-sub currently does support more than 3 layers of playlists/entries. " + "If you encounter this error, please file a ticket with the URLs used." + ) + + mergedeep.merge(kwargs_to_add, self._entry_aggregate_variables()) + for idx, entry_child in enumerate(self.entry_children()): + entry_child.add_kwargs( + self._playlist_variables( + idx=idx, children=self.entry_children(), parent_type=ParentType.PLAYLIST + ) + ) + entry_child.add_kwargs(kwargs_to_add) + + for idx, parent_child in enumerate(self.parent_children()): + parent_child.add_kwargs( + self._playlist_variables( + idx=idx, children=self.parent_children(), parent_type=ParentType.SOURCE + ) + ) + parent_child._set_child_variables(parents=parents + [parent_child]) + + return self + + def _read_children_from_entry_dicts(self, entry_dicts: List[Dict]) -> "EntryParent": """ Populates a tree of EntryParents that belong to this instance """ - child_entries: List["EntryParent"] = [] - - for entry_dict in entry_dicts: - if entry_dict.get("playlist_id") == self.uid: - child_entries.append( - self.__class__( - entry_dict=entry_dict, - working_directory=self.working_directory(), - ) - ) - child_entries[-1].read_children_from_entry_dicts(entry_dicts) + child_entries = [ + EntryParent( + entry_dict=entry_dict, + working_directory=self.working_directory(), + )._read_children_from_entry_dicts(entry_dicts) + for entry_dict in entry_dicts + if entry_dict in self + ] self.child_entries = sorted(child_entries, key=lambda entry: entry.kwargs("playlist_index")) return self - def child_count(self) -> int: - """ - Returns - ------- - Number of child entries - """ - return len(self.child_entries) - def get_thumbnail_url(self, thumbnail_id: str) -> Optional[str]: """ Downloads a specific thumbnail from a YTDL entry's thumbnail list @@ -75,25 +144,94 @@ class EntryParent(BaseEntry): return thumbnail["url"] return None + def __contains__(self, item: Dict | BaseEntry) -> bool: + playlist_id: Optional[str] = None + if isinstance(item, dict): + playlist_id = item.get("playlist_id") + elif isinstance(item, BaseEntry): + playlist_id = item.kwargs_get("playlist_id") + + if not playlist_id: + return False + + return self.uid == playlist_id or any( + child.__contains__(item) for child in self.child_entries + ) + + @classmethod + def _get_disconnected_root_parent( + cls, url: str, parents: List["EntryParent"] + ) -> Optional["EntryParent"]: + """ + Sometimes the root-level parent is disconnected via playlist_ids Find it if it exists. + """ + + def _url_matches(webpage_url: str): + return webpage_url in url or url in webpage_url + + top_level_parents = [ + parent + for parent in parents + if not parent.child_entries and _url_matches(parent.webpage_url) + ] + if len(top_level_parents) == 0: + return None + + match len(top_level_parents): + case 0: + return None + case 1: + return top_level_parents[0] + case _: + raise ValueError( + "Detected multiple top-level parents. " + "Please file an issue on GitHub with the URLs used to produce this error" + ) + @classmethod def from_entry_dicts( - cls, entry_dicts: List[Dict], working_directory: str + cls, url: str, entry_dicts: List[Dict], working_directory: str ) -> List["EntryParent"]: """ Reads all entry dicts and builds a tree of EntryParents """ - return [ + parents = [ EntryParent( entry_dict=entry_dict, working_directory=working_directory - ).read_children_from_entry_dicts(entry_dicts) + )._read_children_from_entry_dicts(entry_dicts) for entry_dict in entry_dicts - if "playlist_id" not in entry_dict + if cls.is_entry_parent(entry_dict) ] - def to_type(self, entry_type: Type[TBaseEntry]) -> TBaseEntry: + if not parents: + return [] + + # If a disconnected root parent exists, connect it here + if (root_parent := cls._get_disconnected_root_parent(url, parents)) is not None: + parents.remove(root_parent) + root_parent.child_entries = parents + parents = [root_parent] + + for parent in parents: + parent._set_child_variables() + + return parents + + # pylint: enable=protected-access + + @classmethod + def from_entry_dicts_with_no_parents( + cls, parents: List["EntryParent"], entry_dicts: List[Dict], working_directory: str + ) -> List[Entry]: """ - Returns - ------- - Converted EntryParent to Entry-like class + Reads all entries that do not have any parents """ - return entry_type(entry_dict=self._kwargs, working_directory=self._working_directory) + + def _in_any_parents(entry_dict: Dict): + return any(entry_dict in parent for parent in parents) + + return [ + Entry(entry_dict=entry_dict, working_directory=working_directory) + for entry_dict in entry_dicts + if cls.is_entry(entry_dict) and not _in_any_parents(entry_dict) + ]