diff --git a/src/ytdl_sub/script/script.py b/src/ytdl_sub/script/script.py index 771f51ee..7873dd13 100644 --- a/src/ytdl_sub/script/script.py +++ b/src/ytdl_sub/script/script.py @@ -9,6 +9,7 @@ from ytdl_sub.script.types.syntax_tree import SyntaxTree from ytdl_sub.script.types.variable import Variable from ytdl_sub.script.utils.exceptions import UNREACHABLE from ytdl_sub.script.utils.exceptions import CycleDetected +from ytdl_sub.script.utils.exceptions import InvalidCustomFunctionArguments from ytdl_sub.script.utils.name_validation import validate_variable_name # pylint: disable=missing-raises-doc @@ -93,6 +94,21 @@ class Script: deps=[], ) + def _ensure_custom_function_arguments_valid(self): + for custom_function_name, custom_function in self._functions.items(): + indices = sorted([arg.index for arg in custom_function.function_arguments]) + if indices != list(range(len(indices))): + if len(indices) == 1: + raise InvalidCustomFunctionArguments( + f"Custom function %{custom_function_name} has invalid function arguments: " + f"The argument must start with $0, not ${indices[0]}." + ) + raise InvalidCustomFunctionArguments( + f"Custom function %{custom_function_name} has invalid function arguments: " + f"{', '.join(sorted(f'${idx}' for idx in indices))} " + f"do not increment from $0 to ${len(indices) - 1}." + ) + def __init__(self, overrides: Dict[str, str]): function_names: Set[str] = { self._function_name(name) for name in overrides.keys() if self._is_function(name) @@ -126,6 +142,7 @@ class Script: } self._ensure_no_custom_function_cycles() + self._ensure_custom_function_arguments_valid() self._ensure_no_variable_cycles() def resolve( diff --git a/src/ytdl_sub/script/types/variable.py b/src/ytdl_sub/script/types/variable.py index 62c68798..1698803c 100644 --- a/src/ytdl_sub/script/types/variable.py +++ b/src/ytdl_sub/script/types/variable.py @@ -13,8 +13,10 @@ class Variable(NamedArgument): class FunctionArgument(Variable): """Arguments for custom functions, i.e. $0, $1, etc""" + index: int + @classmethod def from_idx(cls, idx: int, custom_function_name: Optional[str]) -> "FunctionArgument": if custom_function_name: - return FunctionArgument(name=f"${custom_function_name}___{idx}") - return FunctionArgument(name=f"${idx}") + return FunctionArgument(name=f"${custom_function_name}___{idx}", index=idx) + return FunctionArgument(name=f"${idx}", index=idx) diff --git a/src/ytdl_sub/script/utils/exceptions.py b/src/ytdl_sub/script/utils/exceptions.py index fa67e1f0..dc9e63b6 100644 --- a/src/ytdl_sub/script/utils/exceptions.py +++ b/src/ytdl_sub/script/utils/exceptions.py @@ -19,6 +19,10 @@ class InvalidFunctionName(UserException): """Custom function name is invalid""" +class InvalidCustomFunctionArguments(UserException): + """Custom function arguments are invalid (i.e. they do not increment)""" + + class InvalidCustomFunctionArgumentName(UserException): """Custom function argument name (i.e. $0) is invalid""" diff --git a/tests/unit/script/types/test_custom_function.py b/tests/unit/script/types/test_custom_function.py index 01708f98..649df654 100644 --- a/tests/unit/script/types/test_custom_function.py +++ b/tests/unit/script/types/test_custom_function.py @@ -3,11 +3,11 @@ import re import pytest from ytdl_sub.script.script import Script -from ytdl_sub.script.types.array import ResolvedArray from ytdl_sub.script.types.resolvable import Integer from ytdl_sub.script.utils.exceptions import CycleDetected from ytdl_sub.script.utils.exceptions import FunctionDoesNotExist from ytdl_sub.script.utils.exceptions import InvalidCustomFunctionArgumentName +from ytdl_sub.script.utils.exceptions import InvalidCustomFunctionArguments class TestCustomFunction: @@ -90,7 +90,7 @@ class TestCustomFunction: "$3.14", ], ) - def test_custom_function_invalid_function_arguments(self, name: str): + def test_custom_function_invalid_function_argument_names(self, name: str): with pytest.raises( InvalidCustomFunctionArgumentName, match=re.escape( @@ -104,3 +104,46 @@ class TestCustomFunction: "output": "{%func0(1)}", } ).resolve() + + @pytest.mark.parametrize( + "argument", + [ + "$1", + "$2", + "$3", + ], + ) + def test_custom_function_invalid_function_argument_single(self, argument: str): + with pytest.raises( + InvalidCustomFunctionArguments, + match=re.escape( + f"Custom function %func1 has invalid function arguments: " + f"The argument must start with $0, not {argument}." + ), + ): + Script( + { + "%func1": f"{{[{argument}]}}", + } + ).resolve() + + @pytest.mark.parametrize( + "arguments", + [ + "$0, $2", + "$1, $2, $3", + ], + ) + def test_custom_function_invalid_function_argument_out_of_order(self, arguments: str): + with pytest.raises( + InvalidCustomFunctionArguments, + match=re.escape( + f"Custom function %func1 has invalid function arguments: " + f"{arguments} do not increment from $0 to ${len(arguments.split(',')) - 1}." + ), + ): + Script( + { + "%func1": f"{{[{arguments}]}}", + } + ).resolve()