diff --git a/docs/source/config_reference/scripting/scripting_functions.rst b/docs/source/config_reference/scripting/scripting_functions.rst index 5a02ca39..04dbbb49 100644 --- a/docs/source/config_reference/scripting/scripting_functions.rst +++ b/docs/source/config_reference/scripting/scripting_functions.rst @@ -38,10 +38,11 @@ array_apply_fixed array_at ~~~~~~~~ -:spec: ``array_at(array: Array, idx: Integer) -> AnyArgument`` +:spec: ``array_at(array: Array, idx: Integer, default: Optional[AnyArgument]) -> AnyArgument`` :description: - Return the element in the Array at index ``idx``. + Return the element in the Array at index ``idx``. If ``idx`` exceeds the array length, + either return ``default`` if provided or throw an error. array_contains ~~~~~~~~~~~~~~ diff --git a/src/ytdl_sub/script/functions/array_functions.py b/src/ytdl_sub/script/functions/array_functions.py index 4af6f63c..a07585e7 100644 --- a/src/ytdl_sub/script/functions/array_functions.py +++ b/src/ytdl_sub/script/functions/array_functions.py @@ -73,12 +73,18 @@ class ArrayFunctions: return Array(output) @staticmethod - def array_at(array: Array, idx: Integer) -> AnyArgument: + def array_at(array: Array, idx: Integer, default: Optional[AnyArgument] = None) -> AnyArgument: """ :description: - Return the element in the Array at index ``idx``. + Return the element in the Array at index ``idx``. If ``idx`` exceeds the array length, + either return ``default`` if provided or throw an error. """ - return array.value[idx.value] + try: + return array.value[idx.value] + except IndexError: + if default is not None: + return default + raise @staticmethod def array_first(array: Array, fallback: AnyArgument) -> AnyArgument: diff --git a/tests/unit/script/functions/test_array_functions.py b/tests/unit/script/functions/test_array_functions.py index 06a3b150..bfee05da 100644 --- a/tests/unit/script/functions/test_array_functions.py +++ b/tests/unit/script/functions/test_array_functions.py @@ -26,6 +26,14 @@ class TestArrayFunctions: output = single_variable_output("{%array_at(['a', 'b', 'c'], 1)}") assert output == "b" + def test_array_at_default(self): + output = single_variable_output("{%array_at(['a', 'b', 'c'], 30, 'd')}") + assert output == "d" + + def test_array_at_error(self): + with pytest.raises(FunctionRuntimeException): + single_variable_output("{%array_at(['a', 'b', 'c'], 30)}") + def test_array_flatten(self): output = single_variable_output("{%array_flatten(['a', ['b'], [['c']]])}") assert output == ["a", "b", "c"]