From eb23f59aa02feedef121a05615b2a784af9f9000 Mon Sep 17 00:00:00 2001 From: Jesse Bannon Date: Wed, 13 Dec 2023 17:07:07 -0800 Subject: [PATCH] kwargs not exposed outside of entry --- src/ytdl_sub/downloaders/url/downloader.py | 22 +++------- src/ytdl_sub/entries/base_entry.py | 30 +++++++------ src/ytdl_sub/entries/entry.py | 42 ++++++++++++++++++- src/ytdl_sub/entries/entry_parent.py | 29 +++++++------ src/ytdl_sub/plugins/chapters.py | 4 +- src/ytdl_sub/plugins/split_by_chapters.py | 26 +++--------- src/ytdl_sub/utils/chapters.py | 5 +-- .../enhanced_download_archive.py | 2 +- tests/unit/entries/test_entry.py | 7 ---- 9 files changed, 85 insertions(+), 82 deletions(-) diff --git a/src/ytdl_sub/downloaders/url/downloader.py b/src/ytdl_sub/downloaders/url/downloader.py index ba618532..a0262963 100644 --- a/src/ytdl_sub/downloaders/url/downloader.py +++ b/src/ytdl_sub/downloaders/url/downloader.py @@ -491,21 +491,9 @@ class MultiUrlDownloader(SourcePlugin[MultiUrlValidator]): upload_date_standardized=entry.get(v.upload_date_standardized, str) ) download_idx = self._enhanced_download_archive.num_entries - entry.add( - { - # Tracks number of entries downloaded - v.download_index: download_idx + 1, - # Tracks number of entries with the same upload date to make them unique - v.upload_date_index: upload_date_idx + 1, - v.requested_subtitles: download_entry.kwargs_get( - v.requested_subtitles.metadata_key - ), - v.chapters: download_entry.kwargs_get(v.chapters.metadata_key), - v.sponsorblock_chapters: download_entry.kwargs_get( - v.sponsorblock_chapters.metadata_key - ), - v.comments: download_entry.kwargs_get(v.comments.metadata_key), - } - ) - return entry + return entry.add_injected_variables( + download_entry=download_entry, + download_idx=download_idx, + upload_date_idx=upload_date_idx, + ) diff --git a/src/ytdl_sub/entries/base_entry.py b/src/ytdl_sub/entries/base_entry.py index b3ab6c6b..09a2be1b 100644 --- a/src/ytdl_sub/entries/base_entry.py +++ b/src/ytdl_sub/entries/base_entry.py @@ -42,7 +42,7 @@ class BaseEntry(ABC): str The entry's unique ID """ - return str(self.kwargs(v.uid.metadata_key)) + return str(self._kwargs[v.uid.metadata_key]) @property def extractor(self: "BaseEntry") -> str: @@ -52,21 +52,25 @@ class BaseEntry(ABC): # pylint: disable=line-too-long # Taken from https://github.com/yt-dlp/yt-dlp/blob/e6ab678e36c40ded0aae305bbb866cdab554d417/yt_dlp/YoutubeDL.py#L3514 # pylint: enable=line-too-long - return self.kwargs_get(v.extractor_key.metadata_key) or self.kwargs(v.ie_key.metadata_key) + return ( + self._kwargs_get(v.extractor_key.metadata_key) + or self._kwargs_get(v.ie_key.metadata_key) + or "NO_EXTRACTOR" + ) @property def title(self: "BaseEntry") -> str: """ The title of the entry. If a title does not exist, returns its unique ID. """ - return self.kwargs_get(v.title.metadata_key, self.uid) + return self._kwargs_get(v.title.metadata_key, self.uid) @property def webpage_url(self: "BaseEntry") -> str: """ The url to the webpage. """ - return self.kwargs(v.webpage_url.metadata_key) + return self._kwargs[v.webpage_url.metadata_key] @property def info_json_ext(self) -> str: @@ -78,21 +82,15 @@ class BaseEntry(ABC): """ The uploader id if it exists, otherwise return the unique ID. """ - return self.kwargs_get(v.uploader_id.metadata_key, self.uid) + return self._kwargs_get(v.uploader_id.metadata_key, self.uid) - def kwargs(self, key) -> Any: - """Returns an internal kwarg value supplied from ytdl""" - if key not in self._kwargs: - raise KeyError(f"Expected '{key}' in {self.__class__.__name__} but does not exist.") - return self._kwargs[key] - - def kwargs_get(self, key: str, default: Optional[Any] = None) -> Any: + def _kwargs_get(self, key: str, default: Optional[Any] = None) -> Any: """ Dict get on kwargs """ - if key not in self._kwargs or self.kwargs(key) is None: + if (out := self._kwargs.get(key)) is None: return default - return self.kwargs(key) + return out def working_directory(self) -> str: """ @@ -153,7 +151,7 @@ class BaseEntry(ABC): """ entry_type: Optional[str] = None if isinstance(entry_dict, cls): - entry_type = entry_dict.kwargs_get("_type") + entry_type = entry_dict._kwargs_get("_type") if isinstance(entry_dict, dict): entry_type = entry_dict.get("_type") @@ -168,7 +166,7 @@ class BaseEntry(ABC): """ entry_ext: Optional[str] = None if isinstance(entry_dict, cls): - entry_ext = entry_dict.kwargs_get("ext") + entry_ext = entry_dict._kwargs_get("ext") if isinstance(entry_dict, dict): entry_ext = entry_dict.get("ext") diff --git a/src/ytdl_sub/entries/entry.py b/src/ytdl_sub/entries/entry.py index 90452a23..029b98ed 100644 --- a/src/ytdl_sub/entries/entry.py +++ b/src/ytdl_sub/entries/entry.py @@ -22,6 +22,8 @@ from ytdl_sub.validators.audo_codec_validator import VIDEO_CODEC_EXTS v: VariableDefinitions = VARIABLES _YTDL_SUB_ENTRY_VARIABLES_KWARG_KEY: str = "ytdl_sub_entry_variables" +ytdl_sub_chapters_from_comments = Variable("ytdl_sub_chapters_from_comments") +ytdl_sub_split_by_chapters_parent_uid = Variable("ytdl_sub_split_by_chapters_parent_uid") TypeT = TypeVar("TypeT") @@ -70,6 +72,27 @@ class Entry(BaseEntry, Scriptable): except ScriptVariableNotResolved: return None + def add_injected_variables( + self, download_entry: "Entry", download_idx: int, upload_date_idx: int + ) -> "Entry": + self.add( + { + # Tracks number of entries downloaded + v.download_index: download_idx + 1, + # Tracks number of entries with the same upload date to make them unique + v.upload_date_index: upload_date_idx + 1, + v.requested_subtitles: download_entry._kwargs_get( + v.requested_subtitles.metadata_key, [] + ), + v.chapters: download_entry._kwargs_get(v.chapters.metadata_key, []), + v.sponsorblock_chapters: download_entry._kwargs_get( + v.sponsorblock_chapters.metadata_key, [] + ), + v.comments: download_entry._kwargs_get(v.comments.metadata_key, []), + } + ) + return self + @property def ext(self) -> str: """ @@ -77,7 +100,7 @@ class Entry(BaseEntry, Scriptable): This is not reflected in the entry. See if the mkv file exists and return "mkv" if so, otherwise, return the original extension. """ - ext = self.try_get(v.ext, str) or self.kwargs(key=v.ext.metadata_key) + ext = self.try_get(v.ext, str) or self._kwargs[v.ext.metadata_key] for possible_ext in [ext, "mkv"]: file_path = str(Path(self.working_directory()) / f"{self.uid}.{possible_ext}") if os.path.isfile(file_path): @@ -114,7 +137,7 @@ class Entry(BaseEntry, Scriptable): The source `thumbnail` value and the actual downloaded thumbnail extension sometimes do not match. Return the actual downloaded thumbnail path. """ - thumbnails = self.kwargs_get("thumbnails", []) + thumbnails = self._kwargs_get("thumbnails", []) possible_thumbnail_exts = {"jpg", "webp"} # Always check for jpg and webp thumbs for thumbnail in thumbnails: @@ -195,3 +218,18 @@ class Entry(BaseEntry, Scriptable): Dictionary containing all variables """ return self.script.resolve().as_native() + + @classmethod + def create_split_entry(cls, entry: "Entry", new_uid: str) -> "Entry": + """ + Creates a copy of an entry with a new uid to use as the starting point for a split entry + """ + new_entry = copy.deepcopy(entry) + new_entry._kwargs[v.uid.metadata_key] = new_uid + new_entry.add( + { + v.uid.variable_name: new_uid, + ytdl_sub_split_by_chapters_parent_uid.variable_name: entry.uid, + } + ) + return new_entry diff --git a/src/ytdl_sub/entries/entry_parent.py b/src/ytdl_sub/entries/entry_parent.py index 6c0e2530..4f2669e4 100644 --- a/src/ytdl_sub/entries/entry_parent.py +++ b/src/ytdl_sub/entries/entry_parent.py @@ -19,14 +19,15 @@ v: VariableDefinitions = VARIABLES # pylint: disable=protected-access -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(v.playlist_index.metadata_key, math.inf), ent.uid) - ) - - class EntryParent(BaseEntry): + @classmethod + def _sort_entries(cls, 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(v.playlist_index.metadata_key, math.inf), ent.uid), + ) + def __init__(self, entry_dict: Dict, working_directory: str): super().__init__(entry_dict=entry_dict, working_directory=working_directory) self._parent_children: List["EntryParent"] = [] @@ -57,7 +58,7 @@ class EntryParent(BaseEntry): ) for entry in self.entry_children(): sibling_entry_metadata.append( - {var.metadata_key: entry.kwargs_get(var.metadata_key) for var in variable_filter} + {var.metadata_key: entry._kwargs_get(var.metadata_key) for var in variable_filter} ) return sibling_entry_metadata @@ -79,7 +80,7 @@ class EntryParent(BaseEntry): ) for entry_child in self.entry_children(): - entry_child.add_kwargs(kwargs_to_add) + entry_child._kwargs = dict(entry_child._kwargs, **kwargs_to_add) for parent_child in self.parent_children(): parent_child._set_child_variables(parents=parents + [parent_child]) @@ -99,8 +100,10 @@ class EntryParent(BaseEntry): if entry_dict in self ] - self._parent_children = _sort_entries([ent for ent in entries if self.is_entry_parent(ent)]) - self._entry_children = _sort_entries( + self._parent_children = self._sort_entries( + [ent for ent in entries if self.is_entry_parent(ent)] + ) + self._entry_children = self._sort_entries( [ent.to_type(Entry) for ent in entries if self.is_entry(ent)] ) @@ -119,7 +122,7 @@ class EntryParent(BaseEntry): ------- Desired thumbnail url if it exists. None if it does not. """ - for thumbnail in self.kwargs_get("thumbnails", []): + for thumbnail in self._kwargs_get("thumbnails", []): if thumbnail["id"] == thumbnail_id: return thumbnail["url"] return None @@ -129,7 +132,7 @@ class EntryParent(BaseEntry): if isinstance(item, dict): playlist_id = item.get("playlist_id") elif isinstance(item, BaseEntry): - playlist_id = item.kwargs_get("playlist_id") + playlist_id = item._kwargs_get("playlist_id") if not playlist_id: return False diff --git a/src/ytdl_sub/plugins/chapters.py b/src/ytdl_sub/plugins/chapters.py index 0c28fda3..7a1c9333 100644 --- a/src/ytdl_sub/plugins/chapters.py +++ b/src/ytdl_sub/plugins/chapters.py @@ -10,11 +10,11 @@ from ytdl_sub.config.plugin.plugin_operation import PluginOperation from ytdl_sub.config.validators.options import OptionsDictValidator from ytdl_sub.downloaders.ytdl_options_builder import YTDLOptionsBuilder from ytdl_sub.entries.entry import Entry +from ytdl_sub.entries.entry import ytdl_sub_chapters_from_comments +from ytdl_sub.entries.entry import ytdl_sub_split_by_chapters_parent_uid from ytdl_sub.entries.script.variable_definitions import VARIABLES from ytdl_sub.entries.script.variable_definitions import VariableDefinitions from ytdl_sub.utils.chapters import Chapters -from ytdl_sub.utils.chapters import ytdl_sub_chapters_from_comments -from ytdl_sub.utils.chapters import ytdl_sub_split_by_chapters_parent_uid from ytdl_sub.utils.ffmpeg import set_ffmpeg_metadata_chapters from ytdl_sub.utils.file_handler import FileMetadata from ytdl_sub.validators.regex_validator import RegexListValidator diff --git a/src/ytdl_sub/plugins/split_by_chapters.py b/src/ytdl_sub/plugins/split_by_chapters.py index 50feb0b1..481e8f5c 100644 --- a/src/ytdl_sub/plugins/split_by_chapters.py +++ b/src/ytdl_sub/plugins/split_by_chapters.py @@ -1,5 +1,3 @@ -import copy -from pathlib import Path from typing import Any from typing import Dict from typing import List @@ -11,11 +9,11 @@ from ytdl_sub.config.plugin.plugin import SplitPlugin from ytdl_sub.config.plugin.plugin_operation import PluginOperation from ytdl_sub.config.validators.options import OptionsDictValidator from ytdl_sub.entries.entry import Entry +from ytdl_sub.entries.entry import ytdl_sub_split_by_chapters_parent_uid from ytdl_sub.entries.script.variable_definitions import VARIABLES from ytdl_sub.entries.script.variable_definitions import VariableDefinitions from ytdl_sub.utils.chapters import Chapters from ytdl_sub.utils.chapters import Timestamp -from ytdl_sub.utils.chapters import ytdl_sub_split_by_chapters_parent_uid from ytdl_sub.utils.exceptions import ValidationException from ytdl_sub.utils.ffmpeg import FFMPEG from ytdl_sub.utils.file_handler import FileHandler @@ -188,27 +186,16 @@ class SplitByChaptersPlugin(SplitPlugin[SplitByChaptersOptions]): ) for idx, title in enumerate(chapters.titles): - new_entry = copy.deepcopy(entry) - new_uid = _split_video_uid(source_uid=entry.uid, idx=idx) - - new_entry.add( - { - v.uid.variable_name: new_uid, - ytdl_sub_split_by_chapters_parent_uid.variable_name: entry.uid, - } + new_entry = Entry.create_split_entry( + entry=entry, new_uid=_split_video_uid(source_uid=entry.uid, idx=idx) ) - new_entry.add_kwargs({v.uid.metadata_key: new_uid}) if not self.is_dry_run: - # Get the input/output file paths - input_file = entry.get_download_file_path() - output_file = str(Path(self.working_directory) / f"{new_uid}.{entry.ext}") - # Run ffmpeg to create the split the video FFMPEG.run( _split_video_ffmpeg_cmd( - input_file=input_file, - output_file=output_file, + input_file=entry.get_download_file_path(), + output_file=new_entry.get_download_file_path(), timestamps=chapters.timestamps, idx=idx, ) @@ -219,8 +206,7 @@ class SplitByChaptersPlugin(SplitPlugin[SplitByChaptersOptions]): if entry.is_thumbnail_downloaded(): FileHandler.copy( src_file_path=entry.get_download_thumbnail_path(), - dst_file_path=Path(self.working_directory) - / f"{new_uid}.{entry.get(v.thumbnail_ext, str)}", + dst_file_path=new_entry.get_download_thumbnail_path(), ) # Format the split video diff --git a/src/ytdl_sub/utils/chapters.py b/src/ytdl_sub/utils/chapters.py index 1eacc5b8..9ccfa2c0 100644 --- a/src/ytdl_sub/utils/chapters.py +++ b/src/ytdl_sub/utils/chapters.py @@ -4,16 +4,13 @@ from typing import List from typing import Tuple from ytdl_sub.entries.entry import Entry +from ytdl_sub.entries.entry import ytdl_sub_chapters_from_comments from ytdl_sub.entries.script.variable_definitions import VARIABLES -from ytdl_sub.entries.script.variable_definitions import Variable from ytdl_sub.entries.script.variable_definitions import VariableDefinitions from ytdl_sub.utils.file_handler import FileMetadata v: VariableDefinitions = VARIABLES -ytdl_sub_chapters_from_comments = Variable("ytdl_sub_chapters_from_comments") -ytdl_sub_split_by_chapters_parent_uid = Variable("ytdl_sub_split_by_chapters_parent_uid") - class Timestamp: diff --git a/src/ytdl_sub/ytdl_additions/enhanced_download_archive.py b/src/ytdl_sub/ytdl_additions/enhanced_download_archive.py index eaacbaef..4a86186a 100644 --- a/src/ytdl_sub/ytdl_additions/enhanced_download_archive.py +++ b/src/ytdl_sub/ytdl_additions/enhanced_download_archive.py @@ -14,9 +14,9 @@ from yt_dlp import DateRange from yt_dlp.utils import make_archive_id from ytdl_sub.entries.entry import Entry +from ytdl_sub.entries.entry import ytdl_sub_split_by_chapters_parent_uid from ytdl_sub.entries.script.variable_definitions import VARIABLES from ytdl_sub.entries.script.variable_definitions import VariableDefinitions -from ytdl_sub.utils.chapters import ytdl_sub_split_by_chapters_parent_uid from ytdl_sub.utils.file_handler import FileHandler from ytdl_sub.utils.file_handler import FileHandlerTransactionLog from ytdl_sub.utils.file_handler import FileMetadata diff --git a/tests/unit/entries/test_entry.py b/tests/unit/entries/test_entry.py index df712c4d..4797092d 100644 --- a/tests/unit/entries/test_entry.py +++ b/tests/unit/entries/test_entry.py @@ -15,13 +15,6 @@ class TestEntry(object): for key, expected_value in mock_entry_to_dict.items(): assert out[key] == expected_value, f"{key} does not equal" - def test_entry_missing_kwarg(self, mock_entry): - key = "dne" - expected_error_msg = f"Expected '{key}' in Entry but does not exist." - - with pytest.raises(KeyError, match=expected_error_msg): - mock_entry.kwargs(key) - @pytest.mark.parametrize( "upload_date, year_rev, month_rev, day_rev, month_rev_pad, day_rev_pad", [