[BUGFIX] Do not run all branches of if functions in scripts (#999)

Will only run branch script code of `if` statements if the condition evaluates to that branch

src/ytdl_sub/script/functions/conditional_functions.py

@@ -18,8 +18,8 @@ class ConditionalFunctions:
           depending on the ``condition`` value.
         """
         if condition.value:
-            return true
+            return true.value()
-        return false
+        return false.value()
 
     @staticmethod
     def elif_(*if_elif_else: AnyArgument) -> AnyArgument:
@@ -50,9 +50,9 @@ class ConditionalFunctions:
 
         for idx in range(0, len(arguments) - 1, 2):
             if bool(arguments[idx].value):
-                return arguments[idx + 1]
+                return arguments[idx + 1].value()
 
-        return arguments[-1]
+        return arguments[-1].value()
 
     @staticmethod
     def if_passthrough(
@@ -63,6 +63,7 @@ class ConditionalFunctions:
           Conditional ``if`` statement that returns the ``maybe_true_arg`` if it evaluates to True,
           otherwise returns ``else_arg``.
         """
-        if bool(maybe_true_arg.value):
+        maybe_true_value = maybe_true_arg.value()
-            return maybe_true_arg
+        if bool(maybe_true_value.value):
-        return else_arg
+            return maybe_true_value
+        return else_arg.value()

src/ytdl_sub/script/script.py

@@ -7,6 +7,7 @@ from typing import Set
 from ytdl_sub.script.functions import Functions
 from ytdl_sub.script.parser import parse
 from ytdl_sub.script.script_output import ScriptOutput
+from ytdl_sub.script.types.resolvable import BuiltInFunctionType
 from ytdl_sub.script.types.resolvable import Lambda
 from ytdl_sub.script.types.resolvable import Resolvable
 from ytdl_sub.script.types.syntax_tree import SyntaxTree
@@ -142,22 +143,47 @@ class Script:
                     f"{nested_custom_function.num_input_args}"
                 )
 
+    def _get_lambda_function_names_to_evaluate(self, function: BuiltInFunctionType) -> Set[str]:
+        lambda_function_names: Set[str] = set()
+        for lamb in SyntaxTree(function.args).lambdas:
+            if lamb in function.args:
+                lambda_function_names.add(lamb.value)
+
+            # See if the arg outputs a lambda (from an if).
+            # If so, add the possible lambda to be checked
+            for arg in function.args:
+                if (
+                    isinstance(arg, BuiltInFunctionType)
+                    and arg.output_type() == Lambda
+                    and lamb in arg.args
+                ):
+                    lambda_function_names.add(lamb.value)
+
+        return lambda_function_names
+
     def _ensure_lambda_usage_num_input_arguments_valid(
         self, prefix: str, name: str, definition: SyntaxTree
     ):
         for function in definition.built_in_functions:
-            spec = FunctionSpec.from_callable(Functions.get(function.name))
+            for arg in function.args:
+                self._ensure_lambda_usage_num_input_arguments_valid(
+                    prefix=prefix, name=name, definition=SyntaxTree([arg])
+                )
+
+            spec = FunctionSpec.from_callable(
+                name=function.name, callable_ref=Functions.get(function.name)
+            )
             if not (lambda_type := spec.is_lambda_like):
                 return
 
-            lambda_function_names = set(
+            lambda_function_names = self._get_lambda_function_names_to_evaluate(function=function)
-                lamb.value for lamb in SyntaxTree(function.args).lambdas if isinstance(lamb, Lambda)
-            )
 
             # Only case len(lambda_function_names) > 1 is when used in if-statements
             for lambda_function_name in lambda_function_names:
                 if Functions.is_built_in(lambda_function_name):
-                    lambda_spec = FunctionSpec.from_callable(Functions.get(lambda_function_name))
+                    lambda_spec = FunctionSpec.from_callable(
+                        name=lambda_function_name, callable_ref=Functions.get(lambda_function_name)
+                    )
                     if not lambda_spec.is_num_args_compatible(lambda_type.num_input_args()):
                         expected_args_str = str(lambda_spec.num_required_args)
                         if lambda_spec.num_required_args != len(lambda_spec.args):
@@ -366,13 +392,15 @@ class Script:
 
                 # If the variable's variable dependencies contain an unresolvable variable,
                 # declare it as unresolvable and continue
-                elif definition.contains(unresolvable):
+                elif definition.contains(unresolvable, custom_function_definitions=self._functions):
                     unresolvable.add(variable)
                     del unresolved[variable]
 
                 # Otherwise, if it has dependencies that are all resolved, then
                 # resolve the definition
-                elif not definition.is_subset_of(variables=resolved.keys()):
+                elif not definition.is_subset_of(
+                    variables=resolved.keys(), custom_function_definitions=self._functions
+                ):
                     resolved[variable] = unresolved[variable].resolve(
                         resolved_variables=resolved,
                         custom_functions=self._functions,

src/ytdl_sub/script/types/function.py

@@ -117,7 +117,7 @@ class BuiltInFunction(Function, BuiltInFunctionType):
         -------
         The FunctionSpec of the BuiltInFunction
         """
-        return FunctionSpec.from_callable(self.callable)
+        return FunctionSpec.from_callable(name=self.name, callable_ref=self.callable)
 
     @classmethod
     def _arg_output_type(cls, arg: Argument) -> Type[Argument]:
@@ -229,6 +229,8 @@ class BuiltInFunction(Function, BuiltInFunctionType):
 
         assert isinstance(lambda_array, Array)
 
+        if len(lambda_array.value) == 0:
+            return Array(value=[])
         if len(lambda_array.value) == 1:
             return lambda_array.value[0]
 
@@ -257,14 +259,27 @@ class BuiltInFunction(Function, BuiltInFunctionType):
         resolved_variables: Dict[Variable, Resolvable],
         custom_functions: Dict[str, "VariableDependency"],
     ) -> Resolvable:
+        # TODO: Make conditionals not execute all branches!!!
+        conditional_return_args = self.function_spec.conditional_arg_indices(
+            num_input_args=len(self.args)
+        )
+
         # Resolve all non-lambda arguments
-        resolved_arguments: List[Resolvable | Lambda] = [
+        resolved_arguments: List[Resolvable | Lambda | ReturnableArgument] = [
-            self._resolve_argument_type(
+            (
-                arg=arg,
+                self._resolve_argument_type(
-                resolved_variables=resolved_variables,
+                    arg=arg,
-                custom_functions=custom_functions,
+                    resolved_variables=resolved_variables,
+                    custom_functions=custom_functions,
+                )
+                if idx not in conditional_return_args
+                else ReturnableArgument(
+                    value=functools.partial(
+                        self._resolve_argument_type, arg, resolved_variables, custom_functions
+                    )
+                )
             )
-            for arg in self.args
+            for idx, arg in enumerate(self.args)
         ]
 
         # If a lambda is in a function's arg, resolve it differently

src/ytdl_sub/script/types/variable_dependency.py

@@ -34,17 +34,21 @@ class VariableDependency(ABC):
         Any arguments in the VariableDependency that may or may not need to be resolved.
         """
 
-    def _recurse_get(self, ttype: Type[TypeT], subclass: bool = False) -> List[TypeT]:
+    def _recurse_get(
+        self, ttype: Type[TypeT], subclass: bool = False, instance: bool = True
+    ) -> List[TypeT]:
         output: List[TypeT] = []
         for arg in self._iterable_arguments:
             if subclass and issubclass(type(arg), ttype):
                 output.append(arg)
-            elif isinstance(arg, ttype):
+            elif instance and isinstance(arg, ttype):
+                output.append(arg)
+            elif type(arg) == ttype:  # pylint: disable=unidiomatic-typecheck
                 output.append(arg)
 
             if isinstance(arg, VariableDependency):
                 # pylint: disable=protected-access
-                output.extend(arg._recurse_get(ttype))
+                output.extend(arg._recurse_get(ttype, subclass=subclass, instance=instance))
                 # pylint: enable=protected-access
 
         return output
@@ -57,7 +61,7 @@ class VariableDependency(ABC):
         -------
         All Variables that this depends on.
         """
-        return set(self._recurse_get(Variable))
+        return set(self._recurse_get(Variable, instance=False))
 
     @final
     @property
@@ -156,19 +160,38 @@ class VariableDependency(ABC):
         raise UNREACHABLE
 
     @final
-    def is_subset_of(self, variables: Iterable[Variable]) -> bool:
+    def is_subset_of(
+        self,
+        variables: Iterable[Variable],
+        custom_function_definitions: Dict[str, "VariableDependency"],
+    ) -> bool:
         """
         Returns
         -------
         True if it contains all input variables as a dependency. False otherwise.
         """
+        for custom_function in self.custom_functions:
+            if custom_function_definitions[custom_function.name].is_subset_of(
+                variables=variables, custom_function_definitions=custom_function_definitions
+            ):
+                return True
+
         return not self.variables.issubset(variables)
 
     @final
-    def contains(self, variables: Iterable[Variable]) -> bool:
+    def contains(
+        self,
+        variables: Iterable[Variable],
+        custom_function_definitions: Dict[str, "VariableDependency"],
+    ) -> bool:
         """
         Returns
         -------
         True if it contains any of the input variables. False otherwise.
         """
+        for custom_function in self.custom_functions:
+            if custom_function_definitions[custom_function.name].contains(
+                variables=variables, custom_function_definitions=custom_function_definitions
+            ):
+                return True
         return len(self.variables.intersection(variables)) > 0

src/ytdl_sub/script/utils/type_checking.py

@@ -53,6 +53,25 @@ def get_optional_type(optional_type: Type) -> Type[NamedType]:
     return [arg for arg in optional_type.__args__ if arg != type(None)][0]
 
 
+def _is_union_compatible(
+    arg_type: Type[NamedType],
+    expected_union_type: Type[Resolvable | Optional[Resolvable]],
+) -> bool:
+    if issubclass(arg_type, (NamedCustomFunction, Variable)):
+        return True  # custom-function/variable can be anything, so pass for now
+
+    # if the input arg is a union, do a direct comparison
+    if is_union(arg_type):
+        return arg_type == expected_union_type
+
+    # otherwise, iterate the union to see if it's compatible
+    for union_type in expected_union_type.__args__:
+        if issubclass(arg_type, union_type):
+            return True
+
+    return False
+
+
 def _is_type_compatible(
     arg_type: Type[NamedType],
     expected_arg_type: Type[Resolvable | Optional[Resolvable]],
@@ -63,24 +82,11 @@ def _is_type_compatible(
     True if arg is compatible with expected_arg_type. False otherwise.
     """
     if is_union(expected_arg_type):
-        # See if the arg is a valid against the union
+        return _is_union_compatible(arg_type=arg_type, expected_union_type=expected_arg_type)
-        valid_type = False
 
-        # if the input arg is a union, do a direct comparison
-        if is_union(arg_type):
-            valid_type = arg_type == expected_arg_type
-        # otherwise, iterate the union to see if it's compatible
-        else:
-            for union_type in expected_arg_type.__args__:
-                if issubclass(arg_type, union_type):
-                    valid_type = True
-                    break
-
-        if not valid_type:
-            return False
     # If the input is a union and the expected type is not, see if
     # each possible union input is compatible with the expected type
-    elif is_union(arg_type):
+    if is_union(arg_type):
         for union_type in arg_type.__args__:
             if not _is_type_compatible(union_type, expected_arg_type):
                 return False
@@ -118,6 +124,7 @@ def is_type_compatible(
 
 @dataclass(frozen=True)
 class FunctionSpec:
+    function_name: str
     return_type: Type[Resolvable]
     arg_names: List[str]
     args: Optional[List[Type[Resolvable | Optional[Resolvable]]]] = None
@@ -187,6 +194,22 @@ class FunctionSpec:
             return sum(1 for arg in self.args if not is_optional(arg))
         return 0  # varargs can take any number
 
+    def conditional_arg_indices(self, num_input_args: int) -> List[int]:
+        """
+        Returns
+        -------
+        If the function is conditional, return the indices of the arguments that
+        return for different branches.
+        """
+        if self.function_name == "if":
+            return [1, 2]  # true, false
+        if self.function_name == "elif":
+            # if, retA, elif, retB, retElse
+            return list(range(1, num_input_args, 2)) + [num_input_args - 1]
+        if self.function_name == "if_passthrough":
+            return [0, 1]  # true-passthrough, false-passthrough
+        return []
+
     @property
     def is_lambda_reduce_function(self) -> Optional[Type[LambdaReduce]]:
         """
@@ -261,7 +284,7 @@ class FunctionSpec:
         return self._to_human_readable_name(self.return_type)
 
     @classmethod
-    def from_callable(cls, callable_ref: Callable[..., Resolvable]) -> "FunctionSpec":
+    def from_callable(cls, name: str, callable_ref: Callable[..., Resolvable]) -> "FunctionSpec":
         """
         Returns
         -------
@@ -270,12 +293,14 @@ class FunctionSpec:
         arg_spec: FullArgSpec = inspect.getfullargspec(callable_ref)
         if arg_spec.varargs:
             return FunctionSpec(
+                function_name=name,
                 return_type=arg_spec.annotations["return"],
                 arg_names=[arg_spec.varargs],
                 varargs=arg_spec.annotations[arg_spec.varargs],
             )
 
         return FunctionSpec(
+            function_name=name,
             return_type=arg_spec.annotations["return"],
             arg_names=arg_spec.args,
             args=[arg_spec.annotations[arg_name] for arg_name in arg_spec.args],

src/ytdl_sub/utils/script.py

@@ -13,6 +13,7 @@ from ytdl_sub.script.types.resolvable import Argument
 from ytdl_sub.script.types.resolvable import Boolean
 from ytdl_sub.script.types.resolvable import Float
 from ytdl_sub.script.types.resolvable import Integer
+from ytdl_sub.script.types.resolvable import Lambda
 from ytdl_sub.script.types.resolvable import String
 from ytdl_sub.script.types.variable import Variable
 from ytdl_sub.script.utils.exceptions import UNREACHABLE
@@ -115,6 +116,8 @@ class ScriptUtils:
             out = arg.name
         elif isinstance(arg, Function):
             out = f"%{arg.name}( {', '.join(cls._to_script_code(val) for val in arg.args)} )"
+        elif isinstance(arg, Lambda):
+            out = f"%{arg.value}"
         else:
             raise UNREACHABLE
         return f"{{ {out} }}" if top_level else out

tests/unit/script/functions/test_conditional_functions.py

@@ -101,3 +101,37 @@ class TestConditionalFunction:
                 )
             }"""
             )
+
+    @pytest.mark.parametrize(
+        "function_str, expected_output",
+        [
+            ("{%if(True, True, %assert(False, 'should not reach'))}", True),
+            ("{%if(False, %assert(False, 'should not reach'), False)}", False),
+        ],
+    )
+    def test_if_function_only_evaluates_branch(self, function_str: str, expected_output: bool):
+        output = single_variable_output(function_str)
+        assert output == expected_output
+
+    @pytest.mark.parametrize(
+        "function_str, expected_output",
+        [
+            ("{%elif(True, True, %assert(False, 'should not reach'))}", True),
+            ("{%elif(False, %assert(False, 'should not reach'), False)}", False),
+        ],
+    )
+    def test_elif_function_only_evaluates_branch(self, function_str: str, expected_output: bool):
+        output = single_variable_output(function_str)
+        assert output == expected_output
+
+    @pytest.mark.parametrize(
+        "function_str, expected_output",
+        [
+            ("{%if_passthrough(True, %assert(False, 'should not reach'))}", True),
+        ],
+    )
+    def test_if_passthrough_function_only_evaluates_branch(
+        self, function_str: str, expected_output: bool
+    ):
+        output = single_variable_output(function_str)
+        assert output == expected_output

tests/unit/script/types/test_lambda_function.py

@@ -15,6 +15,11 @@ class TestLambdaFunction:
             {"%times_two": "{%mul($0, 2)}", "wip": "{%array_apply([1, 2, 3], %times_two)}"}
         ).resolve() == ScriptOutput({"wip": Array([Integer(2), Integer(4), Integer(6)])})
 
+    def test_lambda_with_custom_function_empty_input(self):
+        assert Script(
+            {"%times_two": "{%mul($0, 2)}", "wip": "{%array_apply([], %times_two)}"}
+        ).resolve() == ScriptOutput({"wip": Array([])})
+
     def test_conditional_lambda_with_custom_functions(self):
         assert Script(
             {
@@ -54,6 +59,23 @@ class TestLambdaFunction:
             }
         ).resolve() == ScriptOutput({"output": Integer(4)})
 
+    def test_multiple_lambdas_single_definition(self):
+        url_map_def = """{
+                %array_reduce(
+                    %array_apply( array_def, %array_map_format),
+                    %map_extend
+                )
+            }"""
+        script = Script(
+            {
+                "%array_map_format": "{ {$0: $0 } }",
+                "array_def": "{ [1, 2, 3] }",
+                "category_url_map": url_map_def,
+            }
+        )
+
+        assert script.resolve().get("category_url_map").native == {1: 1, 2: 2, 3: 3}
+
 
 class TestLambdaFunctionIncompatibleNumArguments:
     @pytest.mark.parametrize(

tools/docgen/scripting_functions.py

@@ -30,7 +30,7 @@ def function_class_to_name(obj: Type[Any]) -> str:
 
 
 def function_type_hinting(display_function_name: str, function: Any) -> str:
-    spec = FunctionSpec.from_callable(function)
+    spec = FunctionSpec.from_callable(name=display_function_name, callable_ref=function)
     out = ":spec: ``"
     out += display_function_name
     out += spec.human_readable_input_args()