diff --git a/src/ytdl_sub/script/types/function.py b/src/ytdl_sub/script/types/function.py index ab9888b0..705ca3a8 100644 --- a/src/ytdl_sub/script/types/function.py +++ b/src/ytdl_sub/script/types/function.py @@ -58,9 +58,6 @@ class CustomFunction(Function, NamedCustomFunction): resolved_variables: Dict[Variable, Resolvable], custom_functions: Dict[str, "VariableDependency"], ) -> Resolvable: - if NamedCustomFunction(name=self.name) in self.custom_functions: - raise CycleDetected("ackkk!") - resolved_args: List[Resolvable] = [ self._resolve_argument_type( arg=arg, resolved_variables=resolved_variables, custom_functions=custom_functions @@ -74,11 +71,14 @@ class CustomFunction(Function, NamedCustomFunction): resolved_variables_with_args = copy.deepcopy(resolved_variables) for i, arg in enumerate(resolved_args): - function_arg = FunctionArgument.from_idx( - idx=i, custom_function_name=self.name - ) # Function args are 0-based - # if function_arg in resolved_variables_with_args: - # raise StringFormattingException("nested custom functions???") + function_arg = FunctionArgument.from_idx(idx=i, custom_function_name=self.name) + + if function_arg in resolved_variables_with_args: + # function args should always be unique since they are only defined once + # in the custom function as %custom_function_name___idx + # and returned as a set from each custom function. + raise UNREACHABLE + resolved_variables_with_args[function_arg] = arg return custom_functions[self.name].resolve( diff --git a/tests/unit/script/types/test_function.py b/tests/unit/script/types/test_function.py index f4a7a28b..eb560bb5 100644 --- a/tests/unit/script/types/test_function.py +++ b/tests/unit/script/types/test_function.py @@ -118,6 +118,14 @@ class TestFunction: ): Script({"dne": "{%throw}"}).resolve() + def test_custom_function_use_input_param_multiple_times(self): + assert Script( + { + "%custom_square": "{%mul($0, $0)}", + "output": "{%custom_square(3)}", + } + ).resolve() == {"output": Integer(9)} + def test_custom_function_cycle(self): with pytest.raises( CycleDetected, match=re.escape("The custom function %cycle_func cannot call itself.")