diff --git a/src/ytdl_sub/script/script.py b/src/ytdl_sub/script/script.py index 27b4b1e9..f7c20d88 100644 --- a/src/ytdl_sub/script/script.py +++ b/src/ytdl_sub/script/script.py @@ -7,7 +7,8 @@ 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 Lambda, BuiltInFunctionType +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 from ytdl_sub.script.types.variable import FunctionArgument @@ -142,6 +143,24 @@ 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 ): @@ -157,16 +176,7 @@ class Script: if not (lambda_type := spec.is_lambda_like): return - 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) + lambda_function_names = self._get_lambda_function_names_to_evaluate(function=function) # Only case len(lambda_function_names) > 1 is when used in if-statements for lambda_function_name in lambda_function_names: diff --git a/src/ytdl_sub/script/types/variable_dependency.py b/src/ytdl_sub/script/types/variable_dependency.py index ec7be6d5..eb5f646e 100644 --- a/src/ytdl_sub/script/types/variable_dependency.py +++ b/src/ytdl_sub/script/types/variable_dependency.py @@ -43,7 +43,7 @@ class VariableDependency(ABC): output.append(arg) elif instance and isinstance(arg, ttype): output.append(arg) - elif type(arg) == ttype: + elif type(arg) == ttype: # pylint: disable=unidiomatic-typecheck output.append(arg) if isinstance(arg, VariableDependency): diff --git a/src/ytdl_sub/script/utils/type_checking.py b/src/ytdl_sub/script/utils/type_checking.py index 856e75e4..2489760b 100644 --- a/src/ytdl_sub/script/utils/type_checking.py +++ b/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,22 +82,11 @@ def _is_type_compatible( True if arg is compatible with expected_arg_type. False otherwise. """ if is_union(expected_arg_type): - if issubclass(arg_type, (NamedCustomFunction, Variable)): - return True # custom-function/variable can be anything, so pass for now + return _is_union_compatible(arg_type=arg_type, expected_union_type=expected_arg_type) - # if the input arg is a union, do a direct comparison - if is_union(arg_type): - return arg_type == expected_arg_type - - # otherwise, iterate the union to see if it's compatible - for union_type in expected_arg_type.__args__: - if issubclass(arg_type, union_type): - return True - - 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 diff --git a/tests/unit/script/functions/test_conditional_functions.py b/tests/unit/script/functions/test_conditional_functions.py index 5cf04a18..eff73214 100644 --- a/tests/unit/script/functions/test_conditional_functions.py +++ b/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