From e6ee337f20b103914cb94759175eed6a2a5951f2 Mon Sep 17 00:00:00 2001
From: Jesse Bannon <jbann1994@gmail.com>
Date: Mon, 4 Dec 2023 23:21:32 -0800
Subject: [PATCH] register function support

---
 src/ytdl_sub/script/functions/__init__.py | 15 ++++++++++++++-
 src/ytdl_sub/script/types/function.py     | 12 +++++-------
 tests/unit/script/types/test_function.py  | 15 +++++++++++++++
 3 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/src/ytdl_sub/script/functions/__init__.py b/src/ytdl_sub/script/functions/__init__.py
index c0f83784..5e12c931 100644
--- a/src/ytdl_sub/script/functions/__init__.py
+++ b/src/ytdl_sub/script/functions/__init__.py
@@ -1,4 +1,5 @@
 from typing import Callable
+from typing import Dict
 
 from ytdl_sub.script.functions.array_functions import ArrayFunctions
 from ytdl_sub.script.functions.boolean_functions import BooleanFunctions
@@ -22,9 +23,11 @@ class Functions(
     ErrorFunctions,
     RegexFunctions,
 ):
+    _custom_functions: Dict[str, Callable[..., Resolvable]] = {}
+
     @classmethod
     def is_built_in(cls, name: str) -> bool:
-        return hasattr(cls, name) or hasattr(cls, f"{name}_")
+        return hasattr(cls, name) or hasattr(cls, f"{name}_") or name in cls._custom_functions
 
     @classmethod
     def get(cls, name: str) -> Callable[..., Resolvable]:
@@ -32,5 +35,15 @@ class Functions(
             return getattr(cls, name)
         if hasattr(cls, f"{name}_"):
             return getattr(cls, f"{name}_")
+        if name in cls._custom_functions:
+            return cls._custom_functions[name]
 
         raise FunctionDoesNotExistRuntimeException(f"The function {name} does not exist")
+
+    @classmethod
+    def register_function(cls, function: Callable[..., Resolvable]) -> None:
+        if cls.is_built_in(function.__name__):
+            raise ValueError(
+                f"Cannot register a function with name {function.__name__} because it already exists"
+            )
+        cls._custom_functions[function.__name__] = function
diff --git a/src/ytdl_sub/script/types/function.py b/src/ytdl_sub/script/types/function.py
index 8940dab0..566693c1 100644
--- a/src/ytdl_sub/script/types/function.py
+++ b/src/ytdl_sub/script/types/function.py
@@ -91,13 +91,11 @@ class BuiltInFunction(Function, BuiltInFunctionType):
 
     @property
     def callable(self) -> Callable[..., Resolvable]:
-        if hasattr(Functions, self.name):
-            return getattr(Functions, self.name)
-        if hasattr(Functions, self.name + "_"):
-            return getattr(Functions, self.name + "_")
-
-        # Should be validated in the parser
-        raise UNREACHABLE
+        try:
+            return Functions.get(self.name)
+        except Exception as exc:
+            # Should be validated in the parser
+            raise UNREACHABLE from exc
 
     @functools.cached_property
     def function_spec(self) -> FunctionSpec:
diff --git a/tests/unit/script/types/test_function.py b/tests/unit/script/types/test_function.py
index 3ee12b17..372d7d9c 100644
--- a/tests/unit/script/types/test_function.py
+++ b/tests/unit/script/types/test_function.py
@@ -1,9 +1,12 @@
 import re
 
 import pytest
+from unit.script.conftest import single_variable_output
 
+from ytdl_sub.script.functions import Functions
 from ytdl_sub.script.parser import FUNCTION_INVALID_CHAR
 from ytdl_sub.script.script import Script
+from ytdl_sub.script.types.resolvable import Integer
 from ytdl_sub.script.utils.exceptions import FunctionDoesNotExist
 from ytdl_sub.script.utils.exceptions import FunctionRuntimeException
 from ytdl_sub.script.utils.exceptions import IncompatibleFunctionArguments
@@ -14,6 +17,10 @@ def _incompatible_arguments_match(expected: str, recieved: str) -> str:
     return re.escape(f"Expected ({expected})\nReceived ({recieved})")
 
 
+def mock_register_function(integer: Integer) -> Integer:
+    return Integer(integer.value + 100)
+
+
 class TestFunction:
     def test_nested_if_function_incompatible(self):
         function_str = """{
@@ -71,3 +78,11 @@ class TestFunction:
             match=re.escape(str(FUNCTION_INVALID_CHAR)),
         ):
             Script({"dne": "{%throw}"}).resolve()
+
+    def test_register_function(self):
+        try:
+            Functions.register_function(function=mock_register_function)
+            output = single_variable_output(f"{{%mock_register_function(10)}}")
+            assert output == 110
+        finally:
+            del Functions._custom_functions[mock_register_function.__name__]