diff --git a/src/ytdl_sub/script/functions/__init__.py b/src/ytdl_sub/script/functions/__init__.py index 6a1dfc95..00235002 100644 --- a/src/ytdl_sub/script/functions/__init__.py +++ b/src/ytdl_sub/script/functions/__init__.py @@ -1,8 +1,9 @@ from typing import Optional from ytdl_sub.script.functions.numeric_functions import NumericFunctions +from ytdl_sub.script.functions.special_functions import SpecialFunctions from ytdl_sub.script.functions.string_functions import StringFunctions -class Functions(StringFunctions, NumericFunctions): +class Functions(StringFunctions, NumericFunctions, SpecialFunctions): pass diff --git a/src/ytdl_sub/script/functions/special_functions.py b/src/ytdl_sub/script/functions/special_functions.py index 0d390834..8a312de5 100644 --- a/src/ytdl_sub/script/functions/special_functions.py +++ b/src/ytdl_sub/script/functions/special_functions.py @@ -1,10 +1,16 @@ +from typing import Union + from ytdl_sub.script.types.resolvable import Boolean from ytdl_sub.script.types.resolvable import Resolvable +from ytdl_sub.script.types.resolvable import Resolvable_1 +from ytdl_sub.script.types.resolvable import Resolvable_2 class SpecialFunctions: @staticmethod - def if_(condition: Boolean, true: Resolvable, false: Resolvable) -> Resolvable: + def if_( + condition: Boolean, true: Resolvable, false: Resolvable + ) -> Union[Resolvable_1, Resolvable_2]: if condition.value: return true return false diff --git a/src/ytdl_sub/script/parser.py b/src/ytdl_sub/script/parser.py index 8552a485..9516191f 100644 --- a/src/ytdl_sub/script/parser.py +++ b/src/ytdl_sub/script/parser.py @@ -4,7 +4,6 @@ from typing import Optional from ytdl_sub.script.syntax_tree import SyntaxTree from ytdl_sub.script.types.function import ArgumentType from ytdl_sub.script.types.function import Function -from ytdl_sub.script.types.function import IfFunction from ytdl_sub.script.types.resolvable import Boolean from ytdl_sub.script.types.resolvable import Float from ytdl_sub.script.types.resolvable import Integer @@ -158,10 +157,6 @@ class _Parser: while ch := self._read(): if ch == ")": - # Special case for If functions since it can return a Union based on input types - if function_name == "if": - return IfFunction(name=function_name, args=function_args) - return Function(name=function_name, args=function_args) if ch != "(": diff --git a/src/ytdl_sub/script/types/function.py b/src/ytdl_sub/script/types/function.py index 36e29965..c1ccdece 100644 --- a/src/ytdl_sub/script/types/function.py +++ b/src/ytdl_sub/script/types/function.py @@ -20,6 +20,9 @@ from ytdl_sub.script.types.resolvable import Boolean from ytdl_sub.script.types.resolvable import Float from ytdl_sub.script.types.resolvable import Integer from ytdl_sub.script.types.resolvable import Resolvable +from ytdl_sub.script.types.resolvable import Resolvable_0 +from ytdl_sub.script.types.resolvable import Resolvable_1 +from ytdl_sub.script.types.resolvable import Resolvable_2 from ytdl_sub.script.types.resolvable import String from ytdl_sub.script.types.variable import Variable from ytdl_sub.utils.exceptions import StringFormattingException @@ -173,10 +176,12 @@ class Function(VariableDependency): @property def callable(self) -> Callable[..., Resolvable]: - try: + if hasattr(Functions, self.name): return getattr(Functions, self.name) - except AttributeError: - raise StringFormattingException(f"Function name {self.name} does not exist") + if hasattr(Functions, self.name + "_"): + return getattr(Functions, self.name + "_") + + raise StringFormattingException(f"Function name {self.name} does not exist") @functools.cached_property def arg_spec(self) -> FullArgSpec: @@ -188,7 +193,22 @@ class Function(VariableDependency): @property def output_type(self) -> Type[Resolvable]: - return self.arg_spec.annotations["return"] + output_type = self.arg_spec.annotations["return"] + if is_union(output_type): + union_types_list = [] + for union_type in output_type.__args__: + if union_type == Resolvable_0: + union_types_list.append(type(self.args[0])) + elif union_type == Resolvable_1: + union_types_list.append(type(self.args[1])) + elif union_type == Resolvable_2: + union_types_list.append(type(self.args[2])) + else: + union_types_list.append(union_type) + + return Union[tuple(union_types_list)] + + return output_type @property def variables(self) -> Set[Variable]: @@ -208,18 +228,3 @@ class Function(VariableDependency): def resolve(self, resolved_variables: Dict[Variable, Resolvable]) -> Resolvable: raise NotImplemented() - - -@dataclass(frozen=True) -class IfFunction(Function): - def __post_init__(self): - super().__post_init__() - assert len(self.args) == 3 # bool, true, false - - @property - def callable(self) -> Callable[..., Resolvable]: - return SpecialFunctions.if_ - - @property - def output_type(self) -> Type[Resolvable]: - return Union[self.args[1].__class__, self.args[2].__class__] diff --git a/src/ytdl_sub/script/types/resolvable.py b/src/ytdl_sub/script/types/resolvable.py index 9b135506..f7e1b813 100644 --- a/src/ytdl_sub/script/types/resolvable.py +++ b/src/ytdl_sub/script/types/resolvable.py @@ -10,8 +10,20 @@ T = TypeVar("T") NumericT = TypeVar("NumericT", bound=int | float) +class Resolvable_0(ABC): + pass + + +class Resolvable_1(ABC): + pass + + +class Resolvable_2(ABC): + pass + + @dataclass(frozen=True) -class Resolvable(ABC): +class Resolvable(Resolvable_0, Resolvable_1, Resolvable_2, ABC): value: Any def __str__(self) -> str: diff --git a/tests/unit/script/test_parser.py b/tests/unit/script/test_parser.py index c083b33c..e304f09f 100644 --- a/tests/unit/script/test_parser.py +++ b/tests/unit/script/test_parser.py @@ -5,7 +5,6 @@ import pytest from ytdl_sub.script.parser import parse from ytdl_sub.script.syntax_tree import SyntaxTree from ytdl_sub.script.types.function import Function -from ytdl_sub.script.types.function import IfFunction from ytdl_sub.script.types.resolvable import Boolean from ytdl_sub.script.types.resolvable import Float from ytdl_sub.script.types.resolvable import Integer @@ -34,7 +33,7 @@ class TestParser: assert parsed == SyntaxTree( [ String("hello "), - IfFunction( + Function( name="if", args=[Boolean(value=True), String(value="hi"), Float(value=3.4)] ), ] @@ -49,7 +48,7 @@ class TestParser: Function( name="concat", args=[ - IfFunction( + Function( name="if", args=[Boolean(value=True), String("hi"), String("mom")] ), String(value="and dad"), @@ -66,7 +65,7 @@ class TestParser: Function( name="string", args=[ - IfFunction(name="if", args=[Boolean(True), String("hi"), Integer(4)]), + Function(name="if", args=[Boolean(True), String("hi"), Integer(4)]), ], ), ] diff --git a/tests/unit/script/test_syntax_tree.py b/tests/unit/script/test_syntax_tree.py index 9c3a0b0f..6e65bc21 100644 --- a/tests/unit/script/test_syntax_tree.py +++ b/tests/unit/script/test_syntax_tree.py @@ -1,15 +1,8 @@ from typing import Dict -from typing import Union import pytest -from ytdl_sub.script.parser import parse from ytdl_sub.script.syntax_tree import SyntaxTree -from ytdl_sub.script.types.function import Function -from ytdl_sub.script.types.function import IfFunction -from ytdl_sub.script.types.resolvable import Boolean -from ytdl_sub.script.types.resolvable import Float -from ytdl_sub.script.types.resolvable import Integer from ytdl_sub.script.types.resolvable import String from ytdl_sub.script.types.variable import Variable from ytdl_sub.utils.exceptions import StringFormattingException