From 7ac525f8753c91ba4d12f3ef57e8f818dc556564 Mon Sep 17 00:00:00 2001
From: Jesse Bannon <jbann1994@gmail.com>
Date: Tue, 30 Dec 2025 23:46:07 -0800
Subject: [PATCH] [BUGFIX] Fix custom function lambda variable dependency
 issue, part 2 (#1403)

Part two of https://github.com/jmbannon/ytdl-sub/pull/1402 which actually completes the fix.
---
 src/ytdl_sub/script/script.py                 |  8 +++++
 .../script/types/variable_dependency.py       | 29 +++++++++++--------
 .../validators/string_formatter_validators.py |  5 ++++
 .../unit/script/types/test_lambda_function.py |  2 +-
 4 files changed, 31 insertions(+), 13 deletions(-)

diff --git a/src/ytdl_sub/script/script.py b/src/ytdl_sub/script/script.py
index 67eeeaf1..b4c431df 100644
--- a/src/ytdl_sub/script/script.py
+++ b/src/ytdl_sub/script/script.py
@@ -288,6 +288,14 @@ class Script:
                 unresolvable=unresolvable,
             )
 
+        for lambda_func in current_var.lambdas:
+            if lambda_func.value in self._functions:
+                subset_to_resolve |= self._recursive_get_unresolved_output_filter_variables(
+                    current_var=self._functions[lambda_func.value],
+                    subset_to_resolve=subset_to_resolve,
+                    unresolvable=unresolvable,
+                )
+
         return subset_to_resolve
 
     def _get_unresolved_output_filter(
diff --git a/src/ytdl_sub/script/types/variable_dependency.py b/src/ytdl_sub/script/types/variable_dependency.py
index 3e8bfdb4..18f2dab2 100644
--- a/src/ytdl_sub/script/types/variable_dependency.py
+++ b/src/ytdl_sub/script/types/variable_dependency.py
@@ -114,7 +114,6 @@ class VariableDependency(ABC):
                 output.add(ParsedCustomFunction(name=arg.name, num_input_args=len(arg.args)))
             if isinstance(arg, VariableDependency):
                 output.update(arg.custom_functions)
-            # if isinstance(arg, Lambda)
 
         return output
 
@@ -160,6 +159,20 @@ class VariableDependency(ABC):
 
         raise UNREACHABLE
 
+    @final
+    def _custom_function_dependencies(
+        self, custom_function_definitions: Dict[str, "VariableDependency"]
+    ) -> Set[ParsedCustomFunction]:
+        custom_functions = self.custom_functions
+        for lambda_func in self.lambdas:
+            if lambda_func.value in custom_function_definitions:
+                custom_functions.add(
+                    ParsedCustomFunction(
+                        name=lambda_func.value, num_input_args=lambda_func.num_input_args()
+                    )
+                )
+        return custom_functions
+
     @final
     def is_subset_of(
         self,
@@ -171,7 +184,8 @@ class VariableDependency(ABC):
         -------
         True if it contains all input variables as a dependency. False otherwise.
         """
-        for custom_function in self.custom_functions:
+        # If there are lambdas, see if they are custom functions. If so, check them
+        for custom_function in self._custom_function_dependencies(custom_function_definitions):
             if not custom_function_definitions[custom_function.name].is_subset_of(
                 variables=variables, custom_function_definitions=custom_function_definitions
             ):
@@ -191,16 +205,7 @@ class VariableDependency(ABC):
         True if it contains any of the input variables. False otherwise.
         """
         # If there are lambdas, see if they are custom functions. If so, check them
-        custom_functions_to_check = self.custom_functions
-        for lambda_func in self.lambdas:
-            if lambda_func.value in custom_function_definitions:
-                custom_functions_to_check.add(
-                    ParsedCustomFunction(
-                        name=lambda_func.value, num_input_args=lambda_func.num_input_args()
-                    )
-                )
-
-        for custom_function in custom_functions_to_check:
+        for custom_function in self._custom_function_dependencies(custom_function_definitions):
             if custom_function_definitions[custom_function.name].contains(
                 variables=variables, custom_function_definitions=custom_function_definitions
             ):
diff --git a/src/ytdl_sub/validators/string_formatter_validators.py b/src/ytdl_sub/validators/string_formatter_validators.py
index 19a59344..5a44066e 100644
--- a/src/ytdl_sub/validators/string_formatter_validators.py
+++ b/src/ytdl_sub/validators/string_formatter_validators.py
@@ -241,6 +241,11 @@ def _validate_formatter(
     variable_names = {var.name for var in parsed.variables}
     custom_function_names = {f"%{func.name}" for func in parsed.custom_functions}
 
+    # Add lambda functions to custom function names, if it's custom
+    for lambda_func in parsed.lambdas:
+        if lambda_func in mock_script.function_names:
+            custom_function_names.add(lambda_func.value)
+
     if not variable_names.issubset(mock_script.variable_names):
         raise StringFormattingVariableNotFoundException(
             "contains the following variables that do not exist: "
diff --git a/tests/unit/script/types/test_lambda_function.py b/tests/unit/script/types/test_lambda_function.py
index 1112afc4..aac78101 100644
--- a/tests/unit/script/types/test_lambda_function.py
+++ b/tests/unit/script/types/test_lambda_function.py
@@ -81,7 +81,7 @@ class TestLambdaFunction:
             Script(
                 {
                     "the_array": "{ ['a', 'B', 'c', 'D'] }",
-                    "output": "{ %array_apply(the_array, %lower) }",
+                    "output": "{ %array_apply(the_array, %custom_cap) }",
                     "should_lower": "{%bool(True)}",
                     "%custom_cap": """{
                         %if(