From ba71abc4669b05af6522bc035d8c7d87d4f7bda9 Mon Sep 17 00:00:00 2001 From: Jesse Bannon Date: Sat, 6 Jan 2024 23:47:09 -0800 Subject: [PATCH] [BACKEND] Optional default in `%array_at` scripting function (#879) As title --- .../scripting/scripting_functions.rst | 5 +++-- src/ytdl_sub/script/functions/array_functions.py | 12 +++++++++--- tests/unit/script/functions/test_array_functions.py | 8 ++++++++ tests/unit/script/types/test_function.py | 6 +++++- 4 files changed, 25 insertions(+), 6 deletions(-) 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"] diff --git a/tests/unit/script/types/test_function.py b/tests/unit/script/types/test_function.py index 014e4a85..e13828f4 100644 --- a/tests/unit/script/types/test_function.py +++ b/tests/unit/script/types/test_function.py @@ -49,7 +49,11 @@ class TestFunction: @pytest.mark.parametrize( "function_str, expected_types, received_types", [ - ("{%array_at({'a': 'dict?'}, 1)}", "array: Array, idx: Integer", "Map, Integer"), + ( + "{%array_at({'a': 'dict?'}, 1)}", + "array: Array, idx: Integer, default: Optional[AnyArgument]", + "Map, Integer", + ), ("{%array_extend('not', 'array')}", "arrays: Array, ...", "String, String"), ( "{%replace('hi mom', 'mom', 'dad', 1, 0)}",