From df09fd5a9900433991cb348c83a5d16d2af31e82 Mon Sep 17 00:00:00 2001 From: Jesse Bannon Date: Thu, 30 Nov 2023 08:55:48 -0800 Subject: [PATCH] map apply and enumerate --- .../script/functions/map_functions.py | 28 +++++++++++++++++ src/ytdl_sub/script/types/resolvable.py | 11 +++++++ src/ytdl_sub/script/utils/type_checking.py | 11 +++++-- .../script/functions/test_map_functions.py | 30 +++++++++++++++++++ 4 files changed, 78 insertions(+), 2 deletions(-) diff --git a/src/ytdl_sub/script/functions/map_functions.py b/src/ytdl_sub/script/functions/map_functions.py index 6603e923..2746f858 100644 --- a/src/ytdl_sub/script/functions/map_functions.py +++ b/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()) + ] + ) diff --git a/src/ytdl_sub/script/types/resolvable.py b/src/ytdl_sub/script/types/resolvable.py index 11d2b751..d8019c5c 100644 --- a/src/ytdl_sub/script/types/resolvable.py +++ b/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 diff --git a/src/ytdl_sub/script/utils/type_checking.py b/src/ytdl_sub/script/utils/type_checking.py index 24dcb85b..848944cc 100644 --- a/src/ytdl_sub/script/utils/type_checking.py +++ b/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 diff --git a/tests/unit/script/functions/test_map_functions.py b/tests/unit/script/functions/test_map_functions.py index 162d4b77..174824b3 100644 --- a/tests/unit/script/functions/test_map_functions.py +++ b/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"]]