diff --git a/docs/config.rst b/docs/config.rst index ca8aa9809..b1f05b7f2 100644 --- a/docs/config.rst +++ b/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() diff --git a/pyproject.toml b/pyproject.toml index d17d6c0ce..f04e93342 100644 --- a/pyproject.toml +++ b/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 ] diff --git a/src/ytdl_sub/cli/parsers/dl.py b/src/ytdl_sub/cli/parsers/dl.py index 157a156f8..a3ada3b3c 100644 --- a/src/ytdl_sub/cli/parsers/dl.py +++ b/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 @@ def _find_largest_consecutive(cls, indices: List[int]) -> int: 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 @@ def _argument_name_and_value_to_dict(cls, arg_name: str, arg_value: str) -> Dict 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 diff --git a/src/ytdl_sub/config/preset_class_mappings.py b/src/ytdl_sub/config/plugin_mapping.py similarity index 95% rename from src/ytdl_sub/config/preset_class_mappings.py rename to src/ytdl_sub/config/plugin_mapping.py index 31a96363f..5a446a261 100644 --- a/src/ytdl_sub/config/preset_class_mappings.py +++ b/src/ytdl_sub/config/plugin_mapping.py @@ -17,6 +17,7 @@ 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 diff --git a/src/ytdl_sub/config/preset.py b/src/ytdl_sub/config/preset.py index f5ecd3ace..e88e4fc08 100644 --- a/src/ytdl_sub/config/preset.py +++ b/src/ytdl_sub/config/preset.py @@ -12,7 +12,7 @@ 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 diff --git a/src/ytdl_sub/plugins/throttle_protection.py b/src/ytdl_sub/plugins/throttle_protection.py new file mode 100644 index 000000000..25b5791b1 --- /dev/null +++ b/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) diff --git a/src/ytdl_sub/validators/validators.py b/src/ytdl_sub/validators/validators.py index 91b1cf6c7..1220cfca5 100644 --- a/src/ytdl_sub/validators/validators.py +++ b/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 diff --git a/tests/conftest.py b/tests/conftest.py index 185825773..f396e28ce 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -9,6 +9,7 @@ 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 _wrapped_debug(*args, **kwargs): yield for call_args in patched_debug.call_args_list: - if expected_message in call_args.args[0]: - return - - assert False, f"{expected_message} was not found in a logger.debug call" + occurrences += int(expected_message in call_args.args[0]) + + 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: diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 250d439b2..d2957a276 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -9,7 +9,6 @@ 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 diff --git a/tests/e2e/plugins/test_throttle_protection.py b/tests/e2e/plugins/test_throttle_protection.py new file mode 100644 index 000000000..9c52053de --- /dev/null +++ b/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 diff --git a/tests/e2e/youtube/test_video.py b/tests/e2e/youtube/test_video.py index 2b63d6db2..f2c78bb85 100644 --- a/tests/e2e/youtube/test_video.py +++ b/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 @@ def test_single_video_old_video_tags_format_download( 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, diff --git a/tests/resources/expected_downloads_summaries/youtube/test_video.json b/tests/resources/expected_downloads_summaries/youtube/test_video.json index bdf3a2234..4cbb7eca9 100644 --- a/tests/resources/expected_downloads_summaries/youtube/test_video.json +++ b/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" } \ No newline at end of file diff --git a/tests/unit/cli/test_download_args_parser.py b/tests/unit/cli/test_download_args_parser.py index 56e2f8749..675f40663 100644 --- a/tests/unit/cli/test_download_args_parser.py +++ b/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): diff --git a/tests/unit/config/test_config_file.py b/tests/unit/config/test_config_file.py index 6dcb67694..265cf8400 100644 --- a/tests/unit/config/test_config_file.py +++ b/tests/unit/config/test_config_file.py @@ -5,8 +5,8 @@ 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 diff --git a/tests/unit/prebuilt_presets/conftest.py b/tests/unit/conftest.py similarity index 100% rename from tests/unit/prebuilt_presets/conftest.py rename to tests/unit/conftest.py diff --git a/tests/unit/plugins/test_throttle_protection.py b/tests/unit/plugins/test_throttle_protection.py new file mode 100644 index 000000000..0d72e9fd5 --- /dev/null +++ b/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)