From 573c59a7aeeb81d41623406dce7c8c9b0ee37f6e Mon Sep 17 00:00:00 2001 From: Jesse Bannon Date: Wed, 17 May 2023 20:48:19 -0700 Subject: [PATCH] [FEATURE] Perform regex at both metadata and post-download --- src/ytdl_sub/entries/variables/kwargs.py | 1 + src/ytdl_sub/plugins/regex.py | 51 +++++++++++++++++++++++- 2 files changed, 51 insertions(+), 1 deletion(-) diff --git a/src/ytdl_sub/entries/variables/kwargs.py b/src/ytdl_sub/entries/variables/kwargs.py index 5071cd54..9e46062d 100644 --- a/src/ytdl_sub/entries/variables/kwargs.py +++ b/src/ytdl_sub/entries/variables/kwargs.py @@ -46,6 +46,7 @@ UPLOAD_DATE_INDEX = _("upload_date_index", backend=True) REQUESTED_SUBTITLES = _("requested_subtitles", backend=True) CHAPTERS = _("chapters", backend=True) YTDL_SUB_CUSTOM_CHAPTERS = _("ytdl_sub_custom_chapters", backend=True) +YTDL_SUB_REGEX_SOURCE_VARS = _("ytdl_sub_regex_source_vars", backend=True) SPONSORBLOCK_CHAPTERS = _("sponsorblock_chapters", backend=True) SPLIT_BY_CHAPTERS_PARENT_ENTRY = _("split_by_chapters_parent_entry", backend=True) COMMENTS = _("comments", backend=True) diff --git a/src/ytdl_sub/plugins/regex.py b/src/ytdl_sub/plugins/regex.py index 51269083..c9219588 100644 --- a/src/ytdl_sub/plugins/regex.py +++ b/src/ytdl_sub/plugins/regex.py @@ -6,6 +6,7 @@ from typing import Optional from yt_dlp.utils import sanitize_filename from ytdl_sub.entries.entry import Entry +from ytdl_sub.entries.variables.kwargs import YTDL_SUB_REGEX_SOURCE_VARS from ytdl_sub.plugins.plugin import Plugin from ytdl_sub.plugins.plugin import PluginOptions from ytdl_sub.plugins.plugin import PluginPriority @@ -245,12 +246,35 @@ class RegexPlugin(Plugin[RegexOptions]): modify_entry=PluginPriority.MODIFY_ENTRY_AFTER_SPLIT + 0, ) - def modify_entry(self, entry: Entry) -> Optional[Entry]: + @classmethod + def _add_processed_regex_source_var(cls, entry: Entry, source_var: str) -> None: + if not entry.kwargs_contains(YTDL_SUB_REGEX_SOURCE_VARS): + entry.add_kwargs({YTDL_SUB_REGEX_SOURCE_VARS: []}) + + entry.kwargs(YTDL_SUB_REGEX_SOURCE_VARS).append(source_var) + + @classmethod + def _contains_processed_regex_source_var(cls, entry: Entry, source_var: str) -> bool: + return source_var in entry.kwargs_get(YTDL_SUB_REGEX_SOURCE_VARS, []) + + def _record_source_vars_to_process(self, entry: Entry) -> None: + + entry.add_kwargs( + { + YTDL_SUB_REGEX_SOURCE_VARS: list( + self.plugin_options.source_variable_capture_dict.keys() + ) + } + ) + + def _modify_entry_metadata(self, entry: Entry, is_metadata_stage: bool) -> Optional[Entry]: """ Parameters ---------- entry Entry to add source variables to + is_metadata_stage + Whether this is called at the metadata stage or modify stage Returns ------- @@ -265,6 +289,19 @@ class RegexPlugin(Plugin[RegexOptions]): # Iterate each source var to capture and add to the entry for source_var, regex_options in self.plugin_options.source_variable_capture_dict.items(): + + # Record which regex source variables are processed, to + # process as many variables as possible in the metadata stage, then the rest + # after the media file has been downloaded. + if self._contains_processed_regex_source_var(entry, source_var): + continue + + # Continue if the source var isn't in the dict since entry variables could be added + # after the metadata stage + if is_metadata_stage and source_var not in entry_variable_dict: + continue + + self._add_processed_regex_source_var(entry, source_var) maybe_capture = regex_options.match.match_any(input_str=entry_variable_dict[source_var]) # If no capture @@ -331,3 +368,15 @@ class RegexPlugin(Plugin[RegexOptions]): ) return entry + + def modify_entry_metadata(self, entry: Entry) -> Optional[Entry]: + """ + Perform regex at the metadata stage + """ + return self._modify_entry_metadata(entry, is_metadata_stage=True) + + def modify_entry(self, entry: Entry) -> Optional[Entry]: + """ + Perform regex at the metadata stage + """ + return self._modify_entry_metadata(entry, is_metadata_stage=False)