diff --git a/src/ytdl_sub/config/plugin/plugin.py b/src/ytdl_sub/config/plugin/plugin.py index 6b5929c1..0cd02b53 100644 --- a/src/ytdl_sub/config/plugin/plugin.py +++ b/src/ytdl_sub/config/plugin/plugin.py @@ -119,6 +119,17 @@ class Plugin(BasePlugin[OptionsValidatorT], Generic[OptionsValidatorT], ABC): """ return None + def post_completion_entry(self, file_metadata: FileMetadata) -> None: + """ + After the entry file is moved to its final location, run this hook. + + Parameters + ---------- + file_metadata + Metadata about the completed entry's file download + """ + return None + def post_process_subscription(self): """ After all downloaded files have been post-processed, apply a subscription-wide post process diff --git a/src/ytdl_sub/config/plugin/plugin_mapping.py b/src/ytdl_sub/config/plugin/plugin_mapping.py index ac1b851c..bf458d6a 100644 --- a/src/ytdl_sub/config/plugin/plugin_mapping.py +++ b/src/ytdl_sub/config/plugin/plugin_mapping.py @@ -92,6 +92,12 @@ class PluginMapping: EmbedThumbnailPlugin, ] + _ORDER_POST_COMPLETION: List[Type[Plugin]] = [ + # Throttle protection should always be last + # to not sleep over other logic + ThrottleProtectionPlugin + ] + @classmethod def _order_by( cls, plugin_types: List[Type[Plugin]], operation: PluginOperation @@ -102,6 +108,8 @@ class PluginMapping: ordering = cls._ORDER_MODIFY_ENTRY elif operation == PluginOperation.POST_PROCESS: ordering = cls._ORDER_POST_PROCESS + elif operation == PluginOperation.POST_COMPLETION: + ordering = cls._ORDER_POST_COMPLETION else: raise ValueError("PluginOperation does not support ordering") diff --git a/src/ytdl_sub/config/plugin/plugin_operation.py b/src/ytdl_sub/config/plugin/plugin_operation.py index 2c187f30..97706d53 100644 --- a/src/ytdl_sub/config/plugin/plugin_operation.py +++ b/src/ytdl_sub/config/plugin/plugin_operation.py @@ -6,3 +6,4 @@ class PluginOperation(Enum): MODIFY_ENTRY_METADATA = 0 MODIFY_ENTRY = 1 POST_PROCESS = 2 + POST_COMPLETION = 3 diff --git a/src/ytdl_sub/plugins/throttle_protection.py b/src/ytdl_sub/plugins/throttle_protection.py index cf01e537..7acad676 100644 --- a/src/ytdl_sub/plugins/throttle_protection.py +++ b/src/ytdl_sub/plugins/throttle_protection.py @@ -275,6 +275,10 @@ class ThrottleProtectionPlugin(Plugin[ThrottleProtectionOptions]): self._subscription_download_counter: int = 0 self._subscription_max_downloads: Optional[int] = None + # Compute this during post-processing using entry metadata. + # Apply the sleep post-completion. + self._entry_sleep_time: Optional[float] = None + # If subscriptions have a max download limit, set it here for the first subscription if self.plugin_options.max_downloads_per_subscription: self._subscription_max_downloads = ( @@ -333,16 +337,19 @@ class ThrottleProtectionPlugin(Plugin[ThrottleProtectionOptions]): self._subscription_download_counter += 1 if self.plugin_options.sleep_per_download_s: - sleep_time = self.plugin_options.sleep_per_download_s.randomized_float( + self._entry_sleep_time = self.plugin_options.sleep_per_download_s.randomized_float( overrides=self.overrides, entry=entry ) - # pylint: disable=logging-fstring-interpolation) - # needed to test logs in unit test - logger.info(f"Sleeping between downloads for {sleep_time:.2f} seconds") - self.perform_sleep(sleep_time) return None + def post_completion_entry(self, file_metadata: FileMetadata) -> None: + if self._entry_sleep_time: + # pylint: disable=logging-fstring-interpolation) + # needed to test logs in unit test + logger.info(f"Sleeping between downloads for {self._entry_sleep_time:.2f} seconds") + self.perform_sleep(self._entry_sleep_time) + def post_process_subscription(self): # Reset counter to 0 for the next subscription self._subscription_download_counter = 0 diff --git a/src/ytdl_sub/subscriptions/subscription_download.py b/src/ytdl_sub/subscriptions/subscription_download.py index 2fd9c89e..43b811a4 100644 --- a/src/ytdl_sub/subscriptions/subscription_download.py +++ b/src/ytdl_sub/subscriptions/subscription_download.py @@ -259,6 +259,9 @@ class SubscriptionDownload(BaseSubscription, ABC): if self.maintain_download_archive: self.download_archive.save_download_mappings() + for plugin in PluginMapping.order_plugins_by(plugins, PluginOperation.POST_COMPLETION): + plugin.post_completion_entry(file_metadata=entry_metadata) + def _process_entry( self, plugins: List[Plugin], dry_run: bool, entry: Entry, entry_metadata: FileMetadata ) -> None: