From 9b7ba5cb3e63b151fab95a07756b37ca7cd2eef2 Mon Sep 17 00:00:00 2001
From: Jesse Bannon <jbann1994@gmail.com>
Date: Tue, 14 Nov 2023 20:37:20 -0800
Subject: [PATCH] function type checking wip

---
 src/ytdl_sub/script/types/function.py    |  9 +++---
 tests/unit/script/types/test_function.py | 39 +++++++++++++++++++++++-
 2 files changed, 43 insertions(+), 5 deletions(-)

diff --git a/src/ytdl_sub/script/types/function.py b/src/ytdl_sub/script/types/function.py
index 745b8c4b..86b5c7b2 100644
--- a/src/ytdl_sub/script/types/function.py
+++ b/src/ytdl_sub/script/types/function.py
@@ -43,7 +43,7 @@ class FunctionInputSpec:
         input_arg: ArgumentType,
         expected_arg_type: Type[Resolvable | Optional[Resolvable]],
     ) -> bool:
-        if isinstance(input_arg, Function):
+        if isinstance(input_arg, BuiltInFunction):
             input_arg_type = input_arg.output_type
         elif isinstance(input_arg, Variable):
             return True  # unresolved variables can be anything, so pass for now
@@ -115,7 +115,7 @@ class FunctionInputSpec:
             return f"({self.varargs.__name__}, ...)"
 
     @classmethod
-    def from_function(cls, func: "Function") -> "FunctionInputSpec":
+    def from_function(cls, func: "BuiltInFunction") -> "FunctionInputSpec":
         if func.arg_spec.varargs:
             return FunctionInputSpec(varargs=func.arg_spec.annotations[func.arg_spec.varargs])
 
@@ -164,7 +164,7 @@ class Function(VariableDependency, ArgumentType, ABC):
     @classmethod
     def from_name_and_args(cls, name: str, args: List[ArgumentType]) -> "Function":
         if hasattr(Functions, name) or hasattr(Functions, name + "_"):
-            return BuiltInFunction(name=name, args=args)
+            return BuiltInFunction(name=name, args=args).validate_args()
 
         return CustomFunction(name=name, args=args)
 
@@ -214,12 +214,13 @@ class BuiltInFunction(Function):
 
         return f"Expected {self.input_spec.expected_args_str()}.\nReceived {received_args_str}"
 
-    def __post_init__(self):
+    def validate_args(self) -> "BuiltInFunction":
         if not self.input_spec.is_compatible(input_args=self.args):
             raise StringFormattingException(
                 f"Invalid arguments passed to function {self.name}.\n"
                 f"{self._expected_received_error_msg()}"
             )
+        return self
 
     @property
     def callable(self) -> Callable[..., Resolvable]:
diff --git a/tests/unit/script/types/test_function.py b/tests/unit/script/types/test_function.py
index 6e7ded5d..97203cd7 100644
--- a/tests/unit/script/types/test_function.py
+++ b/tests/unit/script/types/test_function.py
@@ -11,6 +11,7 @@ from ytdl_sub.script.parser import UNEXPECTED_COMMA_ARGUMENT
 from ytdl_sub.script.parser import ArgumentParser
 from ytdl_sub.script.script import Script
 from ytdl_sub.script.types.array import ResolvedArray
+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
@@ -18,4 +19,40 @@ from ytdl_sub.script.utils.exceptions import InvalidSyntaxException
 
 
 class TestFunction:
-    pass
+    @pytest.mark.parametrize(
+        "function_str, expected_output",
+        [
+            ("{%if(True, True, False)}", True),
+            ("{%if(False, True, False)}", False),
+        ],
+    )
+    def test_if_function(self, function_str: str, expected_output: bool):
+        assert Script({"func": function_str}).resolve() == {
+            "func": Boolean(expected_output),
+        }
+
+    def test_nested_if_function(self):
+        function_str = """{
+            %if(
+                True,
+                %if(
+                    True,
+                    %if(
+                        True,
+                        "winner",
+                        True
+                    ),
+                    True
+                ),
+                True
+            )
+        }"""
+        assert Script({"func": function_str}).resolve() == {
+            "func": String("winner"),
+        }
+
+    @pytest.mark.parametrize(
+        "function_str", ["{%array_at({'a': 'dict?'}, 1)}" "{%array_extend('not', 'array')}"]
+    )
+    def test_incompatible_types(self, function_str):
+        assert Script({"func": function_str}).resolve()