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..015580dd 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,25 @@ 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 _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 +279,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 +358,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)