diff --git a/src/ytdl_sub/downloaders/downloader.py b/src/ytdl_sub/downloaders/downloader.py index 621554c3..325e6898 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 @@ -439,24 +425,22 @@ class Downloader(DownloadArchiver, Generic[DownloaderOptionsT, DownloaderEntryT] ) return self.parents, orphans - def _download_url( + def _download( self, - collection_url: CollectionUrlValidator, - parents: List[EntryParent], - orphans: List[Entry] = None, + 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: @@ -469,9 +453,7 @@ class Downloader(DownloadArchiver, Generic[DownloaderOptionsT, DownloaderEntryT] # download the bottom-most urls first since they are top-priority for collection_url in reversed(self.collection.collection_urls.list): parents, orphan_entries = self._download_url_metadata(collection_url=collection_url) - for entry in self._download_url( - collection_url=collection_url, parents=parents, orphans=orphan_entries - ): + 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 03162f0c..2db3a6cf 100644 --- a/src/ytdl_sub/downloaders/youtube/channel.py +++ b/src/ytdl_sub/downloaders/youtube/channel.py @@ -148,7 +148,7 @@ class YoutubeChannelDownloader(YoutubeDownloader[YoutubeChannelDownloaderOptions } ) - for entry in super()._download_url(collection_url=collection_url, parents=self.parents): + for entry in super()._download(parents=self.parents): yield entry.to_type(YoutubeVideo) def _download_thumbnail( diff --git a/src/ytdl_sub/downloaders/youtube/playlist.py b/src/ytdl_sub/downloaders/youtube/playlist.py index 87278be1..a3aed608 100644 --- a/src/ytdl_sub/downloaders/youtube/playlist.py +++ b/src/ytdl_sub/downloaders/youtube/playlist.py @@ -113,5 +113,5 @@ class YoutubePlaylistDownloader( } ) - for entry in super()._download_url(collection_url=collection_url, parents=self.parents): + for entry in super()._download(parents=self.parents): yield entry.to_type(YoutubePlaylistVideo) diff --git a/src/ytdl_sub/entries/base_entry.py b/src/ytdl_sub/entries/base_entry.py index 185a276d..b3b1d57d 100644 --- a/src/ytdl_sub/entries/base_entry.py +++ b/src/ytdl_sub/entries/base_entry.py @@ -82,6 +82,20 @@ class BaseEntryVariables: """ return "info.json" + 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, + } + # pylint: enable=no-member @@ -141,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 diff --git a/src/ytdl_sub/entries/entry_parent.py b/src/ytdl_sub/entries/entry_parent.py index 5fd70bbc..d4c9520e 100644 --- a/src/ytdl_sub/entries/entry_parent.py +++ b/src/ytdl_sub/entries/entry_parent.py @@ -1,12 +1,26 @@ +import math import os from typing import Dict from typing import List from typing import Optional +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 +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_id", math.inf), ent.uid)) + + class EntryParent(BaseEntry): def __init__(self, entry_dict: Dict, working_directory: str): super().__init__(entry_dict=entry_dict, working_directory=working_directory) @@ -14,26 +28,57 @@ class EntryParent(BaseEntry): def parent_children(self) -> List["EntryParent"]: """This parent's children that are also parents""" - return [child for child in self.child_entries if self.is_entry_parent(child)] + return _sort_entries([child for child in self.child_entries if self.is_entry_parent(child)]) def entry_children(self) -> List[Entry]: """This parent's children that are entries""" - return [child.to_type(Entry) for child in self.child_entries if self.is_entry(child)] + 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 _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()}, + ) + + # pylint: disable=protected-access + + def _set_child_variables(self, parents: Optional[List["EntryParent"]] = None) -> "EntryParent": + if parents is None: + parents = [self] + + 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." + ) + + for entry_child in self.entry_children(): + entry_child.add_kwargs(kwargs_to_add) + + for parent_child in self.parent_children(): + 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(), - ).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 @@ -80,7 +125,7 @@ class EntryParent(BaseEntry): 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 cls.is_entry_parent(entry_dict) ] @@ -96,10 +141,15 @@ class EntryParent(BaseEntry): parents.remove(first_parent) first_parent.child_entries = parents - return [first_parent] + parents = [first_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