make lambda resolvable

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]:
+    def is_lambda_function(self) -> bool:
-        if Lambda in (self.input_spec.args or []):
+        return Lambda in (self.input_spec.args or [])
-            return [lam for lam in self.args if isinstance(lam, Lambda)][0]
-        return None
 
     @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
+        function_input_lambda_args = [arg for arg in resolved_arguments if isinstance(arg, Lambda)]
-        lambda_function_name = self.lambda_argument.function_name
+        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,
             )

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

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)])
+        }