map apply and enumerate

src/ytdl_sub/script/functions/map_functions.py

@@ -1,9 +1,14 @@
 from typing import Optional
 
+from ytdl_sub.script.types.array import Array
+from ytdl_sub.script.types.array import ResolvedArray
 from ytdl_sub.script.types.map import Map
 from ytdl_sub.script.types.resolvable import AnyArgument
 from ytdl_sub.script.types.resolvable import Boolean
 from ytdl_sub.script.types.resolvable import Hashable
+from ytdl_sub.script.types.resolvable import Integer
+from ytdl_sub.script.types.resolvable import Lambda2
+from ytdl_sub.script.types.resolvable import Lambda3
 from ytdl_sub.script.utils.exceptions import KeyDoesNotExistRuntimeException
 
 
@@ -29,3 +34,26 @@ class MapFunctions:
         Returns True if the key is in the Map. False otherwise.
         """
         return Boolean(key in mapping.value)
+
+    # pylint: disable=unused-argument
+
+    @staticmethod
+    def map_apply(mapping: Map, lambda_function: Lambda2) -> Array:
+        """
+        Apply a lambda function on the Map, where each arg
+        passed to the lambda function is ``key, value`` as two separate args.
+        """
+        return ResolvedArray([ResolvedArray([key, value]) for key, value in mapping.value.items()])
+
+    @staticmethod
+    def map_enumerate(mapping: Map, lambda_function: Lambda3) -> Array:
+        """
+        Apply a lambda function on the Map, where each arg
+        passed to the lambda function is ``idx, key, value`` as three separate args.
+        """
+        return ResolvedArray(
+            [
+                ResolvedArray([Integer(idx), key_value[0], key_value[1]])
+                for idx, key_value in enumerate(mapping.value.items())
+            ]
+        )

src/ytdl_sub/script/types/resolvable.py

@@ -176,3 +176,14 @@ class Lambda2(Lambda):
     @classmethod
     def num_input_args(cls) -> int:
         return 2
+
+
+@dataclass(frozen=True)
+class Lambda3(Lambda):
+    """
+    Type-hinting for functions that apply lambdas with three inputs per element
+    """
+
+    @classmethod
+    def num_input_args(cls) -> int:
+        return 3

src/ytdl_sub/script/utils/type_checking.py

@@ -5,6 +5,7 @@ from typing import Callable
 from typing import List
 from typing import Optional
 from typing import Type
+from typing import TypeVar
 from typing import Union
 from typing import get_origin
 
@@ -13,6 +14,7 @@ from ytdl_sub.script.types.resolvable import BuiltInFunctionType
 from ytdl_sub.script.types.resolvable import FunctionType
 from ytdl_sub.script.types.resolvable import Lambda
 from ytdl_sub.script.types.resolvable import Lambda2
+from ytdl_sub.script.types.resolvable import Lambda3
 from ytdl_sub.script.types.resolvable import NamedType
 from ytdl_sub.script.types.resolvable import Resolvable
 from ytdl_sub.script.types.variable import Variable
@@ -21,6 +23,9 @@ from ytdl_sub.script.utils.exceptions import UNREACHABLE
 # pylint: disable=missing-raises-doc
 
 
+TLambda = TypeVar("TLambda", bound=Lambda)
+
+
 def is_union(arg_type: Type) -> bool:
     """
     Returns
@@ -156,10 +161,12 @@ class FunctionSpec:
         return 0  # varargs can take any number
 
     @property
-    def is_lambda_function(self) -> Optional[Type[Lambda | Lambda2]]:
+    def is_lambda_function(self) -> Optional[Type[TLambda]]:
+        if Lambda3 in (self.args or []):
+            return Lambda3
         if Lambda2 in (self.args or []):
             return Lambda2
-        elif Lambda in (self.args or []):
+        if Lambda in (self.args or []):
             return Lambda
         return None
 

tests/unit/script/functions/test_map_functions.py

@@ -68,3 +68,33 @@ class TestMapFunctions:
             .native
         )
         assert output == expected_value
+
+    def test_map_apply(self):
+        output = (
+            Script(
+                {
+                    "%custom_func": "{[%upper($0), %lower($1)]}",
+                    "map1": "{{'Key1': 'Value1', 'Key2': 'Value2'}}",
+                    "output": "{%map_apply(map1, %custom_func)}",
+                }
+            )
+            .resolve(update=True)
+            .get("output")
+            .native
+        )
+        assert output == [["KEY1", "value1"], ["KEY2", "value2"]]
+
+    def test_map_enumerate(self):
+        output = (
+            Script(
+                {
+                    "%custom_func": "{[$0, %upper($1), %lower($2)]}",
+                    "map1": "{{'Key1': 'Value1', 'Key2': 'Value2'}}",
+                    "output": "{%map_enumerate(map1, %custom_func)}",
+                }
+            )
+            .resolve(update=True)
+            .get("output")
+            .native
+        )
+        assert output == [[0, "KEY1", "value1"], [1, "KEY2", "value2"]]