diff --git a/src/ytdl_sub/script/functions/numeric_functions.py b/src/ytdl_sub/script/functions/numeric_functions.py index ba86bde2..beee8c25 100644 --- a/src/ytdl_sub/script/functions/numeric_functions.py +++ b/src/ytdl_sub/script/functions/numeric_functions.py @@ -1,6 +1,7 @@ from ytdl_sub.script.types.resolvable import Float from ytdl_sub.script.types.resolvable import Integer from ytdl_sub.script.types.resolvable import Numeric +from ytdl_sub.script.types.resolvable import Resolvable def _to_numeric(value: int | float) -> Numeric: @@ -10,6 +11,14 @@ def _to_numeric(value: int | float) -> Numeric: class NumericFunctions: + @staticmethod + def float(value: Resolvable) -> Float: + return Float(value=float(value.value)) + + @staticmethod + def int(value: Resolvable) -> Integer: + return Integer(value=int(value.value)) + @staticmethod def add(left: Numeric, right: Numeric) -> Numeric: return _to_numeric(left.value + right.value) diff --git a/src/ytdl_sub/script/functions/string_functions.py b/src/ytdl_sub/script/functions/string_functions.py index 6773bcb4..5d0f68ef 100644 --- a/src/ytdl_sub/script/functions/string_functions.py +++ b/src/ytdl_sub/script/functions/string_functions.py @@ -1,10 +1,15 @@ from typing import Optional from ytdl_sub.script.types.resolvable import Integer +from ytdl_sub.script.types.resolvable import Resolvable from ytdl_sub.script.types.resolvable import String class StringFunctions: + @staticmethod + def string(value: Resolvable) -> String: + return String(value=str(value.value)) + @staticmethod def lower(string: String) -> String: """ diff --git a/src/ytdl_sub/script/types/function.py b/src/ytdl_sub/script/types/function.py index 47aaa7ae..61522a81 100644 --- a/src/ytdl_sub/script/types/function.py +++ b/src/ytdl_sub/script/types/function.py @@ -76,13 +76,26 @@ class FunctionInputSpec: if is_union(expected_arg_type): # See if the arg is a valid against the union valid_type = False - for union_type in expected_arg_type.__args__: - if issubclass(input_arg_type, union_type): - valid_type = True - break + + # if the input arg is a union, do a direct comparison + if is_union(input_arg_type): + valid_type = input_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(input_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(input_arg_type): + for union_type in input_arg_type.__args__: + if not issubclass(union_type, expected_arg_type): + return False + elif not issubclass(input_arg_type, expected_arg_type): return False diff --git a/tests/unit/script/test_parser.py b/tests/unit/script/test_parser.py index f57c9795..5d9bd7da 100644 --- a/tests/unit/script/test_parser.py +++ b/tests/unit/script/test_parser.py @@ -41,7 +41,7 @@ class TestParser: ) assert parsed.ast[1].output_type == Union[String, Float] - def test_conditional_as_input(self): + def test_conditional_as_input_same_outputs(self): parsed = parse("hello {%concat(%if(True, 'hi', 'mom'), 'and dad')}") assert parsed == SyntaxTree( [ @@ -58,6 +58,20 @@ class TestParser: ] ) + def test_conditional_as_input_different_outputs(self): + parsed = parse("hello {%string(%if(True, 'hi', 4))}") + assert parsed == SyntaxTree( + [ + String("hello "), + Function( + name="string", + args=[ + IfFunction(name="if", args=[Boolean(True), String("hi"), Integer(4)]), + ], + ), + ] + ) + def test_single_function_one_vararg(self): parsed = parse("hello {%concat('hi mom')}") assert parsed == SyntaxTree(