From f74ef8227577c7eb1399c2f8dae57f149ef11725 Mon Sep 17 00:00:00 2001 From: Jesse Bannon Date: Fri, 24 Nov 2023 00:37:49 -0800 Subject: [PATCH] getting wild --- src/ytdl_sub/script/script.py | 12 +++++++-- .../script/types/variable_dependency.py | 12 +++++++-- src/ytdl_sub/script/utils/type_checking.py | 12 ++++----- .../unit/script/types/test_lambda_function.py | 26 +++++++++++++------ 4 files changed, 44 insertions(+), 18 deletions(-) diff --git a/src/ytdl_sub/script/script.py b/src/ytdl_sub/script/script.py index b70caed5..6583a3d4 100644 --- a/src/ytdl_sub/script/script.py +++ b/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 diff --git a/src/ytdl_sub/script/types/variable_dependency.py b/src/ytdl_sub/script/types/variable_dependency.py index 8bda071c..4e399c9f 100644 --- a/src/ytdl_sub/script/types/variable_dependency.py +++ b/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]: diff --git a/src/ytdl_sub/script/utils/type_checking.py b/src/ytdl_sub/script/utils/type_checking.py index 3cf6aa69..24dcb85b 100644 --- a/src/ytdl_sub/script/utils/type_checking.py +++ b/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 diff --git a/tests/unit/script/types/test_lambda_function.py b/tests/unit/script/types/test_lambda_function.py index d564fad6..631287a9 100644 --- a/tests/unit/script/types/test_lambda_function.py +++ b/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})}}", } )