From 31d8c0a52f34b83fe762b6ab36d7b68922fb18b2 Mon Sep 17 00:00:00 2001
From: Jesse Bannon <jbann1994@gmail.com>
Date: Tue, 21 Nov 2023 16:56:03 -0800
Subject: [PATCH] make lambda resolvable

---
 src/ytdl_sub/script/types/function.py    | 20 ++++++++++----------
 src/ytdl_sub/script/types/resolvable.py  |  2 +-
 tests/unit/script/types/test_function.py | 11 +++++++++++
 3 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/src/ytdl_sub/script/types/function.py b/src/ytdl_sub/script/types/function.py
index f8be8f0a..b8c1bb11 100644
--- a/src/ytdl_sub/script/types/function.py
+++ b/src/ytdl_sub/script/types/function.py
@@ -28,7 +28,7 @@ from ytdl_sub.script.types.variable import FunctionArgument
 from ytdl_sub.script.types.variable import Variable
 from ytdl_sub.script.types.variable_dependency import VariableDependency
 from ytdl_sub.script.utils.exception_formatters import FunctionArgumentsExceptionFormatter
-from ytdl_sub.script.utils.exceptions import FunctionDoesNotExist
+from ytdl_sub.script.utils.exceptions import FunctionDoesNotExist, UNREACHABLE
 from ytdl_sub.script.utils.exceptions import FunctionRuntimeException
 from ytdl_sub.script.utils.exceptions import UserThrownRuntimeError
 from ytdl_sub.script.utils.type_checking import FunctionInputSpec
@@ -144,10 +144,8 @@ class BuiltInFunction(Function, TypeHintedFunctionType):
         )
 
     @property
-    def lambda_argument(self) -> Optional[Lambda]:
-        if Lambda in (self.input_spec.args or []):
-            return [lam for lam in self.args if isinstance(lam, Lambda)][0]
-        return None
+    def is_lambda_function(self) -> bool:
+        return Lambda in (self.input_spec.args or [])
 
     @classmethod
     def _arg_output_type(cls, arg: ArgumentType) -> Type[ArgumentType]:
@@ -184,8 +182,11 @@ class BuiltInFunction(Function, TypeHintedFunctionType):
             2. Preemptively creating the lambda's unresolved output array using output args from (1)
             3. Resolve it like any other syntax
         """
-        assert self.lambda_argument is not None
-        lambda_function_name = self.lambda_argument.function_name
+        function_input_lambda_args = [arg for arg in resolved_arguments if isinstance(arg, Lambda)]
+        if not self.is_lambda_function or len(function_input_lambda_args) != 1:
+            raise UNREACHABLE
+
+        lambda_function_name = function_input_lambda_args[0].function_name
 
         try:
             lambda_args = self.callable(*resolved_arguments)
@@ -222,13 +223,12 @@ class BuiltInFunction(Function, TypeHintedFunctionType):
                 custom_functions=custom_functions,
             )
             for arg in self.args
-            if not isinstance(arg, Lambda)
         ]
 
         # If a lambda is in a function's arg, resolve it differently
-        if lambda_argument := self.lambda_argument:
+        if self.is_lambda_function:
             return self._resolve_lambda_function(
-                resolved_arguments=resolved_arguments + [lambda_argument],
+                resolved_arguments=resolved_arguments,
                 resolved_variables=resolved_variables,
                 custom_functions=custom_functions,
             )
diff --git a/src/ytdl_sub/script/types/resolvable.py b/src/ytdl_sub/script/types/resolvable.py
index d41d3c0d..0c6d835f 100644
--- a/src/ytdl_sub/script/types/resolvable.py
+++ b/src/ytdl_sub/script/types/resolvable.py
@@ -135,5 +135,5 @@ class TypeHintedFunctionType(FunctionType, ABC):
 
 
 @dataclass(frozen=True)
-class Lambda(ArgumentType):
+class Lambda(Resolvable):
     function_name: str
diff --git a/tests/unit/script/types/test_function.py b/tests/unit/script/types/test_function.py
index 0151af38..02efaae2 100644
--- a/tests/unit/script/types/test_function.py
+++ b/tests/unit/script/types/test_function.py
@@ -121,3 +121,14 @@ class TestFunction:
         assert Script(
             {"%times_two": "{%mul($0, 2)}", "wip": "{%array_apply([1, 2, 3], %times_two)}"}
         ).resolve() == {"wip": ResolvedArray([Integer(2), Integer(4), Integer(6)])}
+
+    def test_conditional_lambda_function(self):
+        assert Script(
+            {
+                "%times_three": "{%mul($0, 3)}",
+                "%times_two": "{%mul($0, 2)}",
+                "wip": "{%array_apply([1, 2, 3], %if(False, %times_two, %times_three))}"
+            }
+        ).resolve() == {
+            "wip": ResolvedArray([Integer(3), Integer(6), Integer(9)])
+        }
\ No newline at end of file