diff --git a/src/ytdl_sub/script/functions/array_functions.py b/src/ytdl_sub/script/functions/array_functions.py index 3a565888..ba610e68 100644 --- a/src/ytdl_sub/script/functions/array_functions.py +++ b/src/ytdl_sub/script/functions/array_functions.py @@ -1,3 +1,4 @@ +import itertools from typing import List from typing import Optional @@ -85,6 +86,17 @@ class ArrayFunctions: """ return Array(list(reversed(array.value))) + @staticmethod + def array_product(*arrays: Array) -> Array: + """ + Returns the Cartesian product of elements from different arrays + """ + out: List[Resolvable] = [] + for combo in itertools.product(*[arr.value for arr in arrays]): + out.append(Array(combo)) + + return Array(out) + # pylint: disable=unused-argument @staticmethod diff --git a/tests/unit/script/functions/test_array_functions.py b/tests/unit/script/functions/test_array_functions.py index 8d2cb29c..a85e38ce 100644 --- a/tests/unit/script/functions/test_array_functions.py +++ b/tests/unit/script/functions/test_array_functions.py @@ -44,6 +44,10 @@ class TestArrayFunctions: output = single_variable_output("{%array_reverse(['a', 'b', 'c'])}") assert output == ["c", "b", "a"] + def test_array_product(self): + output = single_variable_output("{%array_product(['a', 'b', 'c'], ['arg'])}") + assert output == [["a", "arg"], ["b", "arg"], ["c", "arg"]] + def test_array_apply(self): output = single_variable_output("{%array_apply(['a', 'b', 'c'], %capitalize)}") assert output == ["A", "B", "C"]