if now uses psuedo generics

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

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

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 != "(":

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__]

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:

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)]),
                     ],
                 ),
             ]

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