utils: simple function overloading helper

diff --git a/manager/knot_resolver_manager/utils/overload.py b/manager/knot_resolver_manager/utils/overload.py
new file mode 100644 (file)
index 0000000..2b2b039
--- /dev/null
+++ b/manager/knot_resolver_manager/utils/overload.py
@@ -0,0 +1,30 @@
+from typing import Any, Callable, Dict, Generic, Tuple, Type, TypeVar, overload
+
+T = TypeVar("T")
+
+class OverloadedFunctionException(Exception): pass
+
+class overloaded(Generic[T]):
+    def __init__(self):
+        self.vtable: Dict[Tuple[Any], Callable[..., T]] = {}
+    
+    def add(self, *args: Type[Any], **kwargs: Type[Any]) -> Callable[[Callable[..., T]], Callable[..., T]]:
+        if len(kwargs) != 0:
+            raise OverloadedFunctionException("Sorry, named arguments are not supported. You can however implement them and make them functional... ;)")
+
+        def wrapper(func: Callable[...,T]) -> Callable[...,T]:
+            signature = tuple(args)
+            self.vtable[signature] = func
+            def inner_wrapper(*args: Any, **kwargs: Any) -> T:
+                return self(*args, **kwargs)
+            return inner_wrapper
+        return wrapper
+    
+    def __call__(self, *args: Any, **kwargs: Any) -> T:
+        if len(kwargs) != 0:
+            raise OverloadedFunctionException("Sorry, named arguments are not supported. You can however implement them and make them functional... ;)")
+
+        signature = tuple(type(a) for a in args)
+        if signature not in self.vtable:
+            raise OverloadedFunctionException(f"Function overload with signature {signature} is not registered and can't be called.")
+        return self.vtable[signature](*args)
\ No newline at end of file

diff --git a/manager/tests/utils/test_overloaded.py b/manager/tests/utils/test_overloaded.py
new file mode 100644 (file)
index 0000000..283f79e
--- /dev/null
+++ b/manager/tests/utils/test_overloaded.py
@@ -0,0 +1,20 @@
+from knot_resolver_manager.utils.overload import overloaded
+
+
+def test_simple():
+    func: overloaded[None] = overloaded()
+
+    @func.add(int)
+    def f1(a: int) -> None:
+        assert type(a) == int
+    
+    @func.add(str)
+    def f2(a: str) -> None:
+        assert type(a) == str
+    
+    func("test")
+    func(5)
+    f1("test")
+    f2(5)
+    f1("test")
+    f2(5)
\ No newline at end of file