[FEATURE] Add throttle_protection plugin (#799)

Provides options to make ytdl-sub look more 'human-like' to protect from throttling. For range-based values, a random number will be chosen within the range to avoid sleeps looking scripted.

Usage:
```
    throttle_protection:
      sleep_per_download_s:
        min: 2.2
        max: 10.8
      sleep_per_subscription_s:
        min: 9.0
        max: 14.1
      max_downloads_per_subscription:
        min: 10
        max: 36
      subscription_download_probability: 1.0
```

docs/config.rst

@@ -252,6 +252,14 @@ subtitles
 
 -------------------------------------------------------------------------------
 
+throttle_protection
+'''''''''''''''''''
+.. autoclass:: ytdl_sub.plugins.throttle_protection.ThrottleProtectionOptions()
+  :members:
+  :member-order: bysource
+
+-------------------------------------------------------------------------------
+
 video_tags
 ''''''''''
 .. autoclass:: ytdl_sub.plugins.video_tags.VideoTagsOptions()

pyproject.toml

@@ -15,6 +15,7 @@ disable = [
     "R0913", # Too many arguments
     "R0901", # too-many-ancestors
     "R0902", # too-many-instance-attributes
+    "R1711", # useless-return
     "W0511", # TODO
 ]
 

src/ytdl_sub/cli/parsers/dl.py

@@ -1,6 +1,7 @@
 import hashlib
 import re
 import shlex
+from typing import Any
 from typing import Dict
 from typing import List
 from typing import Tuple
@@ -116,7 +117,7 @@ class DownloadArgsParser:
         return largest_consecutive + 1
 
     @classmethod
-    def _argument_name_and_value_to_dict(cls, arg_name: str, arg_value: str) -> Dict:
+    def _argument_name_and_value_to_dict(cls, arg_name: str, arg_value: Any) -> Dict:
         """
         :param arg_name: Argument name in the form of 'key1.key2.key3'
         :param arg_value: Argument value
@@ -134,11 +135,15 @@ class DownloadArgsParser:
 
         next_dict[arg_name_split[-1]] = arg_value
 
-        # TODO: handle ints/floats
-        if arg_value == "True":
-            next_dict[arg_name_split[-1]] = True
-        elif arg_value == "False":
-            next_dict[arg_name_split[-1]] = False
+        if isinstance(arg_value, str):
+            if arg_value == "True":
+                next_dict[arg_name_split[-1]] = True
+            elif arg_value == "False":
+                next_dict[arg_name_split[-1]] = False
+            elif arg_value.isdigit():
+                next_dict[arg_name_split[-1]] = int(arg_value)
+            elif arg_value.replace(".", "", 1).isdigit():
+                next_dict[arg_name_split[-1]] = float(arg_value)
 
         return argument_dict
 

src/ytdl_sub/config/plugin_mapping.py

@@ -17,6 +17,7 @@ from ytdl_sub.plugins.output_directory_nfo_tags import OutputDirectoryNfoTagsPlu
 from ytdl_sub.plugins.regex import RegexPlugin
 from ytdl_sub.plugins.split_by_chapters import SplitByChaptersPlugin
 from ytdl_sub.plugins.subtitles import SubtitlesPlugin
+from ytdl_sub.plugins.throttle_protection import ThrottleProtectionPlugin
 from ytdl_sub.plugins.video_tags import VideoTagsPlugin
 
 
@@ -41,6 +42,7 @@ class PluginMapping:
         "subtitles": SubtitlesPlugin,
         "chapters": ChaptersPlugin,
         "split_by_chapters": SplitByChaptersPlugin,
+        "throttle_protection": ThrottleProtectionPlugin,
     }
 
     @classmethod

src/ytdl_sub/config/preset.py

@@ -12,7 +12,7 @@ from mergedeep import mergedeep
 
 from ytdl_sub.config.config_validator import ConfigValidator
 from ytdl_sub.config.plugin import Plugin
-from ytdl_sub.config.preset_class_mappings import PluginMapping
+from ytdl_sub.config.plugin_mapping import PluginMapping
 from ytdl_sub.config.preset_options import OptionsValidator
 from ytdl_sub.config.preset_options import OutputOptions
 from ytdl_sub.config.preset_options import Overrides

src/ytdl_sub/plugins/throttle_protection.py

@@ -0,0 +1,233 @@
+import random
+import time
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+from ytdl_sub.config.plugin import Plugin
+from ytdl_sub.config.preset_options import OptionsDictValidator
+from ytdl_sub.config.preset_options import Overrides
+from ytdl_sub.entries.entry import Entry
+from ytdl_sub.utils.file_handler import FileMetadata
+from ytdl_sub.utils.logger import Logger
+from ytdl_sub.validators.strict_dict_validator import StrictDictValidator
+from ytdl_sub.validators.validators import FloatValidator
+from ytdl_sub.validators.validators import ProbabilityValidator
+from ytdl_sub.ytdl_additions.enhanced_download_archive import EnhancedDownloadArchive
+
+logger = Logger.get("throttle-protection")
+
+
+class RandomizedRangeValidator(StrictDictValidator):
+    """
+    Validator to specify a float range between [min, max)
+    """
+
+    _required_keys = {"max"}
+    _optional_keys = {"min"}
+
+    def __init__(self, name, value):
+        super().__init__(name, value)
+
+        self._max = self._validate_key(key="max", validator=FloatValidator).value
+        self._min = self._validate_key_if_present(
+            key="min", validator=FloatValidator, default=0.0
+        ).value
+
+        if self._min < 0:
+            raise self._validation_exception("min must be greater than zero")
+
+        if self._max < self._min:
+            raise self._validation_exception(
+                f"max ({self._max}) must be greater than or equal to min ({self._min})"
+            )
+
+    def randomized_float(self) -> float:
+        """
+        Returns
+        -------
+        A random float within the range
+        """
+        return random.uniform(self._min, self._max)
+
+    def randomized_int(self) -> int:
+        """
+        Returns
+        -------
+        A random float within the range, then cast to an integer (floored)
+        """
+        return int(self.randomized_float())
+
+
+class ThrottleProtectionOptions(OptionsDictValidator):
+    """
+    Provides options to make ytdl-sub look more 'human-like' to protect from throttling. For
+    range-based values, a random number will be chosen within the range to avoid sleeps looking
+    scripted.
+
+    Usage:
+
+    .. code-block:: yaml
+
+       presets:
+         my_example_preset:
+           throttle_protection:
+             sleep_per_download_s:
+               min: 2.2
+               max: 10.8
+             sleep_per_subscription_s:
+               min: 9.0
+               max: 14.1
+             max_downloads_per_subscription:
+               min: 10
+               max: 36
+             subscription_download_probability: 1.0
+    """
+
+    _optional_keys = {
+        "sleep_per_download_s",
+        "sleep_per_subscription_s",
+        "max_downloads_per_subscription",
+        "subscription_download_probability",
+    }
+
+    def __init__(self, name, value):
+        super().__init__(name, value)
+
+        self._sleep_per_download_s = self._validate_key_if_present(
+            key="sleep_per_download_s", validator=RandomizedRangeValidator
+        )
+        self._sleep_per_subscription_s = self._validate_key_if_present(
+            key="sleep_per_subscription_s", validator=RandomizedRangeValidator
+        )
+        self._max_downloads_per_subscription = self._validate_key_if_present(
+            key="max_downloads_per_subscription", validator=RandomizedRangeValidator
+        )
+        self._subscription_download_probability = self._validate_key_if_present(
+            key="subscription_download_probability", validator=ProbabilityValidator
+        )
+
+    @property
+    def sleep_per_download_s(self) -> Optional[RandomizedRangeValidator]:
+        """
+        Number in seconds to sleep between each download. Does not include time it takes for
+        ytdl-sub to perform post-processing.
+        """
+        return self._sleep_per_download_s
+
+    @property
+    def sleep_per_subscription_s(self) -> Optional[RandomizedRangeValidator]:
+        """
+        Number in seconds to sleep between each subscription.
+        """
+        return self._sleep_per_subscription_s
+
+    @property
+    def max_downloads_per_subscription(self) -> Optional[RandomizedRangeValidator]:
+        """
+        Number of downloads to perform per subscription.
+        """
+        return self._max_downloads_per_subscription
+
+    @property
+    def subscription_download_probability(self) -> Optional[ProbabilityValidator]:
+        """
+        Probability to perform any downloads, recomputed for each subscription. This is only
+        recommended to set if you run ytdl-sub in a cron-job, that way you are statistically
+        guaranteed over time to eventually download the subscription.
+        """
+        return self._subscription_download_probability
+
+
+class ThrottleProtectionPlugin(Plugin[ThrottleProtectionOptions]):
+    plugin_options_type = ThrottleProtectionOptions
+
+    def __init__(
+        self,
+        options: ThrottleProtectionOptions,
+        overrides: Overrides,
+        enhanced_download_archive: EnhancedDownloadArchive,
+    ):
+        super().__init__(options, overrides, enhanced_download_archive)
+        self._subscription_download_counter: int = 0
+        self._subscription_max_downloads: Optional[int] = 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 = (
+                self.plugin_options.max_downloads_per_subscription.randomized_int()
+            )
+
+    def ytdl_options_match_filters(self) -> Tuple[List[str], List[str]]:
+        """
+        Returns
+        -------
+        If subscription_download_probability, match-filters that will perform no downloads
+        if it's rolled to not download.
+        """
+        perform_download: Tuple[List[str], List[str]] = [], []
+        do_not_perform_download: Tuple[List[str], List[str]] = [], [
+            "title = __YTDL_SUB_THROTTLE_PROTECTION_ON_SUBSCRIPTION_DOWNLOAD__"
+        ]
+
+        if self.plugin_options.subscription_download_probability:
+            proba = self.plugin_options.subscription_download_probability.value
+            # assume proba is set to 1.0, random.random() will always be < 1, can never reach this
+            if random.random() > proba:
+                logger.info(
+                    "Subscription download probability of %f missed, skipping this subscription",
+                    proba,
+                )
+                return do_not_perform_download
+
+        return perform_download
+
+    def modify_entry_metadata(self, entry: Entry) -> Optional[Entry]:
+        if (
+            self._subscription_max_downloads is not None
+            and self._subscription_download_counter >= self._subscription_max_downloads
+        ):
+            if self._subscription_download_counter == self._subscription_max_downloads:
+                logger.info(
+                    "Reached subscription max downloads of %d for throttle protection",
+                    self._subscription_max_downloads,
+                )
+                self._subscription_download_counter += 1  # increment to only print once
+
+            return None
+
+        return entry
+
+    def post_process_entry(self, entry: Entry) -> Optional[FileMetadata]:
+        if (
+            self._subscription_max_downloads is not None
+            and self._subscription_download_counter == 0
+        ):
+            logger.debug(
+                "Setting subscription max downloads to %d", self._subscription_max_downloads
+            )
+
+        # Increment the counter
+        self._subscription_download_counter += 1
+
+        if self.plugin_options.sleep_per_download_s:
+            sleep_time = self.plugin_options.sleep_per_download_s.randomized_float()
+            logger.debug("Sleeping between downloads for %0.2f seconds", sleep_time)
+            time.sleep(sleep_time)
+
+        return None
+
+    def post_process_subscription(self):
+        # Reset counter to 0 for the next subscription
+        self._subscription_download_counter = 0
+
+        # If present, reset max downloads for the next subscription
+        if self.plugin_options.max_downloads_per_subscription:
+            self._subscription_max_downloads = (
+                self.plugin_options.max_downloads_per_subscription.randomized_int
+            )
+
+        if self.plugin_options.sleep_per_subscription_s:
+            sleep_time = self.plugin_options.sleep_per_subscription_s.randomized_float()
+            logger.debug("Sleeping between subscriptions for %0.2f seconds", sleep_time)
+            time.sleep(sleep_time)

src/ytdl_sub/validators/validators.py

@@ -149,6 +149,15 @@ class IntValidator(ValueValidator[int]):
     _expected_value_type_name = "int"
 
 
+class ProbabilityValidator(FloatValidator):
+    _expected_value_type_name = "probability"
+
+    def __init__(self, name, value):
+        super().__init__(name, value)
+        if self.value < 0 or self.value > 1:
+            raise self._validation_exception("Probabilities must be between 0 and 1.0")
+
+
 class ListValidator(Validator, ABC, Generic[ValidatorT]):
     """
     Validates a list of objects to validate

tests/conftest.py

@@ -9,6 +9,7 @@ from typing import Any
 from typing import Callable
 from typing import Dict
 from typing import List
+from typing import Optional
 from unittest.mock import patch
 
 import pytest
@@ -78,11 +79,17 @@ def reformat_directory() -> Path:
 
 
 @contextlib.contextmanager
-def assert_logs(logger: logging.Logger, expected_message: str, log_level: str = "debug"):
+def assert_logs(
+    logger: logging.Logger,
+    expected_message: str,
+    log_level: str = "debug",
+    expected_occurrences: Optional[int] = None,
+):
     """
     Patches any function, but calls the original function.
     Intended to see if the particular function is called.
     """
+    occurrences = 0
     debug_logger = Logger.get()
 
     def _wrapped_debug(*args, **kwargs):
@@ -92,10 +99,14 @@ def assert_logs(logger: logging.Logger, expected_message: str, log_level: str =
         yield
 
     for call_args in patched_debug.call_args_list:
-        if expected_message in call_args.args[0]:
-            return
+        occurrences += int(expected_message in call_args.args[0])
 
-    assert False, f"{expected_message} was not found in a logger.debug call"
+    if expected_occurrences:
+        assert (
+            occurrences == expected_occurrences
+        ), f"{expected_message} was expected {expected_occurrences} times, got {occurrences}"
+    else:
+        assert occurrences > 0, f"{expected_message} was not found in a logger.debug call"
 
 
 def preset_dict_to_dl_args(preset_dict: Dict) -> str:

tests/e2e/conftest.py

@@ -9,7 +9,6 @@ import pytest
 from ytdl_sub.cli.entrypoint import main
 from ytdl_sub.subscriptions.subscription import Subscription
 from ytdl_sub.utils.file_handler import FileHandler
-from ytdl_sub.utils.system import IS_WINDOWS
 
 
 @pytest.fixture

tests/e2e/plugins/test_throttle_protection.py

@@ -0,0 +1,72 @@
+import pytest
+from conftest import assert_logs
+
+from ytdl_sub.plugins.throttle_protection import logger as throttle_protection_logger
+from ytdl_sub.subscriptions.subscription import Subscription
+
+
+@pytest.fixture
+def preset_dict_max_downloads_0(output_directory):
+    return {
+        "preset": "Jellyfin Music Videos",
+        "download": "https://youtube.com/watch?v=HKTNxEqsN3Q",
+        "format": "worst[ext=mp4]",
+        "overrides": {
+            "music_video_artist": "JMC",
+            "music_video_directory": output_directory,
+        },
+        "throttle_protection": {"max_downloads_per_subscription": {"max": 0}},
+    }
+
+
+@pytest.fixture
+def preset_dict_subscription_download_proba_0(output_directory):
+    return {
+        "preset": "Jellyfin Music Videos",
+        "download": "https://youtube.com/watch?v=HKTNxEqsN3Q",
+        "format": "worst[ext=mp4]",
+        "overrides": {
+            "music_video_artist": "JMC",
+            "music_video_directory": output_directory,
+        },
+        "throttle_protection": {"subscription_download_probability": 0.0},
+    }
+
+
+class TestThrottleProtection:
+    def test_max_downloads(
+        self,
+        default_config,
+        preset_dict_max_downloads_0,
+        output_directory,
+    ):
+        single_video_subscription = Subscription.from_dict(
+            config=default_config,
+            preset_name="music_video_single_video_test",
+            preset_dict=preset_dict_max_downloads_0,
+        )
+
+        with assert_logs(
+            logger=throttle_protection_logger,
+            expected_message="Reached subscription max downloads of %d",
+            log_level="info",
+            expected_occurrences=1,
+        ):
+            transaction_log = single_video_subscription.download(dry_run=True)
+
+        assert transaction_log.is_empty
+
+    def test_subscription_probability(
+        self,
+        default_config,
+        preset_dict_subscription_download_proba_0,
+        output_directory,
+    ):
+        single_video_subscription = Subscription.from_dict(
+            config=default_config,
+            preset_name="music_video_single_video_test",
+            preset_dict=preset_dict_subscription_download_proba_0,
+        )
+
+        transaction_log = single_video_subscription.download(dry_run=True)
+        assert transaction_log.is_empty

tests/e2e/youtube/test_video.py

@@ -21,7 +21,6 @@ def single_video_preset_dict_old_video_tags_format(output_directory):
         "download": "https://youtube.com/watch?v=HKTNxEqsN3Q",
         # override the output directory with our fixture-generated dir
         "output_options": {
-            "output_directory": output_directory,
             "maintain_download_archive": False,
         },
         "embed_thumbnail": True,  # embed thumb into the video
@@ -32,7 +31,10 @@ def single_video_preset_dict_old_video_tags_format(output_directory):
                 "title": "{title}",
             }
         },
-        "overrides": {"music_video_artist": "JMC"},
+        "overrides": {
+            "music_video_artist": "JMC",
+            "music_video_directory": output_directory,
+        },
     }
 
 
@@ -43,7 +45,6 @@ def single_video_preset_dict(output_directory):
         "download": "https://youtube.com/watch?v=HKTNxEqsN3Q",
         # override the output directory with our fixture-generated dir
         "output_options": {
-            "output_directory": output_directory,
             "maintain_download_archive": False,
         },
         "embed_thumbnail": True,  # embed thumb into the video
@@ -52,7 +53,12 @@ def single_video_preset_dict(output_directory):
         "video_tags": {
             "title": "{title}",
         },
-        "overrides": {"music_video_artist": "JMC"},
+        # And test subscription download proba = 1.0
+        "throttle_protection": {"subscription_download_probability": 1.0},
+        "overrides": {
+            "music_video_artist": "JMC",
+            "music_video_directory": output_directory,
+        },
     }
 
 
@@ -113,7 +119,7 @@ class TestYoutubeVideo:
             transaction_log_summary_file_name="youtube/test_video.txt",
         )
 
-    @pytest.mark.parametrize("dry_run", [True])
+    @pytest.mark.parametrize("dry_run", [True, False])
     def test_single_video_download(
         self,
         default_config,

tests/resources/expected_downloads_summaries/youtube/test_video.json

@@ -1,6 +1,5 @@
 {
-  "JMC/JMC - Oblivion Mod ＂Falcor＂ p.1-thumb.jpg": "fb95b510681676e81c321171fc23143e",
-  "JMC/JMC - Oblivion Mod ＂Falcor＂ p.1.info.json": "08b0f7d93488d625bd7311ab925cec77",
-  "JMC/JMC - Oblivion Mod ＂Falcor＂ p.1.mp4": "797b44f3207be01651780d6d86cb70bb",
-  "JMC/JMC - Oblivion Mod ＂Falcor＂ p.1.nfo": "24cc4e17d2bebc89b2759ce5471d403e"
+  "JMC/Oblivion Mod ＂Falcor＂ p.1.jpg": "fb95b510681676e81c321171fc23143e",
+  "JMC/Oblivion Mod ＂Falcor＂ p.1.mp4": "0448c9fd3eeaba4eca7f650fb93fe21b",
+  "JMC/Oblivion Mod ＂Falcor＂ p.1.nfo": "58c2be339869b5d071c1758d55c72ddb"
 }

tests/unit/cli/test_download_args_parser.py

@@ -97,6 +97,26 @@ class TestDownloadArgsParser:
                 "dl --parameter.not.using.list[0] 'v0'",
                 {"parameter": {"not": {"using": {"list[0]": "v0"}}}},
             ),
+            (
+                None,
+                "dl --a.float.parameter 1.3",
+                {"a": {"float": {"parameter": 1.3}}},
+            ),
+            (
+                None,
+                "dl --a.int.parameter 6",
+                {"a": {"int": {"parameter": 6}}},
+            ),
+            (
+                None,
+                "dl --a.true.parameter True",
+                {"a": {"true": {"parameter": True}}},
+            ),
+            (
+                None,
+                "dl --a.false.parameter False",
+                {"a": {"false": {"parameter": False}}},
+            ),
         ],
     )
     def test_successful_args(self, config_options_generator, aliases, cmd, expected_sub_dict):

tests/unit/config/test_config_file.py

@@ -5,8 +5,8 @@ from typing import Optional
 import pytest
 
 from ytdl_sub.config.config_file import ConfigFile
+from ytdl_sub.config.plugin_mapping import PluginMapping
 from ytdl_sub.config.preset import PRESET_KEYS
-from ytdl_sub.config.preset_class_mappings import PluginMapping
 from ytdl_sub.utils.exceptions import ValidationException
 
 

tests/unit/plugins/test_throttle_protection.py

@@ -0,0 +1,59 @@
+from conftest import assert_logs
+
+from ytdl_sub.plugins.throttle_protection import logger as throttle_protection_logger
+from ytdl_sub.subscriptions.subscription import Subscription
+
+
+class TestThrottleProtectionPlugin:
+    def test_sleeps_log(
+        self,
+        config,
+        subscription_name,
+        output_directory,
+        mock_download_collection_entries,
+    ):
+        preset_dict = {
+            "preset": [
+                "Kodi Music Videos",
+            ],
+            "overrides": {
+                "url": "https://your.name.here",
+                "music_video_directory": output_directory,
+            },
+            "throttle_protection": {
+                "sleep_per_download_s": {
+                    "min": 0.01,
+                    "max": 0.01,
+                },
+                "sleep_per_subscription_s": {
+                    "min": 0.02,
+                    "max": 0.02,
+                },
+            },
+        }
+
+        subscription = Subscription.from_dict(
+            config=config,
+            preset_name=subscription_name,
+            preset_dict=preset_dict,
+        )
+
+        with mock_download_collection_entries(
+            is_youtube_channel=False, num_urls=1, is_extracted_audio=False
+        ), assert_logs(
+            logger=throttle_protection_logger,
+            expected_message="Sleeping between downloads for %0.2f seconds",
+            log_level="debug",
+            expected_occurrences=4,
+        ):
+            _ = subscription.download(dry_run=False)
+
+        with mock_download_collection_entries(
+            is_youtube_channel=False, num_urls=1, is_extracted_audio=False
+        ), assert_logs(
+            logger=throttle_protection_logger,
+            expected_message="Sleeping between subscriptions for %0.2f seconds",
+            log_level="debug",
+            expected_occurrences=1,
+        ):
+            _ = subscription.download(dry_run=False)