danvics/ytdl-sub
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

[BACKEND] Add source and playlist variables to every entry (#236)

Browse source
This commit is contained in:
Jesse Bannon 2022-09-16 22:56:56 -07:00 committed by GitHub
parent 3719b14103
commit c66e5b72f3
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
8 changed files with 333 additions and 106 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
src/ytdl_sub/config/preset_class_mappings.py
View file

@ -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

75
src/ytdl_sub/downloaders/downloader.py
View file

@ -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):

9
src/ytdl_sub/downloaders/youtube/channel.py
View file

@ -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,

11
src/ytdl_sub/downloaders/youtube/playlist.py
View file

@ -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)

4
src/ytdl_sub/downloaders/youtube/video.py
View file

@ -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)

96
src/ytdl_sub/entries/base_entry.py
View file

@ -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

16
src/ytdl_sub/entries/entry.py
View file

@ -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

224
src/ytdl_sub/entries/entry_parent.py
View file

@ -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)
]