getting wild

src/ytdl_sub/script/script.py

@@ -156,7 +156,11 @@ class Script:
                 if lambda_type := spec.is_lambda_function:
 
                     lambda_function_names = set(
-                        [lamb.value for lamb in function.args if isinstance(lamb, Lambda)]
+                        [
+                            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
@@ -202,7 +206,11 @@ class Script:
                 if lambda_type := spec.is_lambda_function:
 
                     lambda_function_names = set(
-                        [lamb.value for lamb in function.args if isinstance(lamb, Lambda)]
+                        [
+                            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

src/ytdl_sub/script/types/variable_dependency.py

@@ -31,11 +31,14 @@ class VariableDependency(ABC):
     def _iterable_arguments(self) -> List[Argument]:
         pass
 
-    def _recurse_get(self, ttype: Type[TType]) -> List[TType]:
+    def _recurse_get(self, ttype: Type[TType], subclass: bool = False) -> List[TType]:
         output: List[TType] = []
         for arg in self._iterable_arguments:
-            if isinstance(arg, ttype):
+            if subclass and issubclass(type(arg), ttype):
                 output.append(arg)
+            elif isinstance(arg, ttype):
+                output.append(arg)
+
             if isinstance(arg, VariableDependency):
                 output.extend(arg._recurse_get(ttype))
 
@@ -56,6 +59,11 @@ class VariableDependency(ABC):
     def function_arguments(self) -> Set[FunctionArgument]:
         return set(self._recurse_get(FunctionArgument))
 
+    @final
+    @property
+    def lambdas(self) -> Set[Lambda]:
+        return set(self._recurse_get(Lambda, subclass=True))
+
     @final
     @property
     def custom_functions(self) -> Set[ParsedCustomFunction]:

src/ytdl_sub/script/utils/type_checking.py

@@ -60,12 +60,6 @@ def is_type_compatible(
     arg_type: Type[NamedType] = arg.__class__
     if isinstance(arg, BuiltInFunctionType):
         arg_type = arg.output_type()  # built-in function
-    elif isinstance(arg, Lambda):
-        # lambda, check if expected_arg_type is a subclass
-        # Do not return on just that to also allow lambdas to be returned as
-        # ReturnableArguments (i.e in an %if statement)
-        if issubclass(expected_arg_type, arg_type):
-            return True
     elif isinstance(arg, FunctionType):
         return True  # custom-function, can be anything, so pass for now
     elif isinstance(arg, Variable):
@@ -94,6 +88,12 @@ def is_type_compatible(
             if not issubclass(union_type, expected_arg_type):
                 return False
 
+    elif issubclass(arg_type, Lambda) and issubclass(expected_arg_type, arg_type):
+        # lambda, check if expected_arg_type is a subclass
+        # Do not return on just that to also allow lambdas to be returned as
+        # ReturnableArguments (i.e in an %if statement)
+        return True
+
     elif not issubclass(arg_type, expected_arg_type):
         return False
 

tests/unit/script/types/test_lambda_function.py

@@ -53,7 +53,12 @@ class TestLambdaFunction:
             }
         ).resolve() == {"output": Integer(4)}
 
-    def test_custom_function_lambda_in_variable_incompatible_number_of_args(self):
+
+class TestLambdaFunctionIncompatibleNumArguments:
+    @pytest.mark.parametrize(
+        "lambda_value", ["%enumerate_output", "%if(False, %capitalize, %enumerate_output)"]
+    )
+    def test_custom_function_lambda_in_variable(self, lambda_value: str):
         with pytest.raises(
             IncompatibleFunctionArguments,
             match=re.escape(
@@ -65,11 +70,12 @@ class TestLambdaFunction:
                 {
                     "%enumerate_output": "{[$0, $1]}",
                     "array1": "{['a', 'b', 'c']}",
-                    "output": "{%array_apply(array1, %enumerate_output)}",
+                    "output": f"{{%array_apply(array1, {lambda_value})}}",
                 }
             )
 
-    def test_function_lambda_in_variable_incompatible_number_of_args(self):
+    @pytest.mark.parametrize("lambda_value", ["%replace", "%if(False, %capitalize, %replace)"])
+    def test_function_lambda_in_variable(self, lambda_value: str):
         with pytest.raises(
             IncompatibleFunctionArguments,
             match=re.escape(
@@ -80,11 +86,14 @@ class TestLambdaFunction:
             Script(
                 {
                     "array1": "{['a', 'b', 'c']}",
-                    "output": "{%array_apply(array1, %replace)}",
+                    "output": f"{{%array_apply(array1, {lambda_value})}}",
                 }
             )
 
-    def test_custom_function_lambda_in_custom_function_incompatible_number_of_args(self):
+    @pytest.mark.parametrize(
+        "lambda_value", ["%enumerate_output", "%if(False, %concat, %enumerate_output)"]
+    )
+    def test_custom_function_lambda_in_custom_function(self, lambda_value: str):
         with pytest.raises(
             IncompatibleFunctionArguments,
             match=re.escape(
@@ -96,11 +105,12 @@ class TestLambdaFunction:
                 {
                     "%enumerate_output": "{[$0, $1, $2]}",
                     "array1": "{['a', 'b', 'c']}",
-                    "%output": "{%array_enumerate(array1, %enumerate_output)}",
+                    "%output": f"{{%array_enumerate(array1, {lambda_value})}}",
                 }
             )
 
-    def test_function_lambda_in_custom_function_incompatible_number_of_args(self):
+    @pytest.mark.parametrize("lambda_value", ["%replace", "%if(False, %concat, %replace)"])
+    def test_function_lambda_in_custom_function(self, lambda_value: str):
         with pytest.raises(
             IncompatibleFunctionArguments,
             match=re.escape(
@@ -111,6 +121,6 @@ class TestLambdaFunction:
             Script(
                 {
                     "array1": "{['a', 'b', 'c']}",
-                    "%output": "{%array_enumerate(array1, %replace)}",
+                    "%output": f"{{%array_enumerate(array1, {lambda_value})}}",
                 }
             )