draft: first pass at overloading verb_dispatch

.github/workflows/ci.yml

@@ -33,7 +33,7 @@ jobs:
             python-version: 3.8
             requirements: numpy~=1.22.0 pandas~=1.3.5 SQLAlchemy~=1.4.29 psycopg2-binary~=2.9.3 PyMySQL==1.0.2
           - name: "2022-early dependencies"
-            python-version: 3.10.1
+            python-version: "3.10"
             requirements: numpy~=1.22.0 pandas~=1.3.5 SQLAlchemy~=1.4.29 psycopg2-binary~=2.9.3 PyMySQL==1.0.2
             latest: true
 

siuba/siu/dispatchers.py

@@ -1,12 +1,19 @@
 # symbolic dispatch wrapper ---------------------------------------------------
 
-from functools import singledispatch, update_wrapper, wraps
-import inspect
+from functools import singledispatch, wraps, partial
 
-from .calls import Call, FuncArg, MetaArg, Lazy, PipeCall, _Isolate
+from .calls import Call, FuncArg, MetaArg, PipeCall, _Isolate
 from .symbolic import Symbolic, create_sym_call, strip_symbolic
 
-from typing import Callable
+from typing import Callable, overload, TYPE_CHECKING, Any
+
+
+if TYPE_CHECKING:
+    from functools import _SingleDispatchCallable
+
+    _DispatchDecorator = Callable[[Callable[..., Any]], _SingleDispatchCallable[Any]]
+    _DispatchType = type[Any] | tuple[type[Any]]
+
 
 def _dispatch_not_impl(func_name):
     def f(x, *args, **kwargs):
@@ -123,6 +130,12 @@ def f_dispatch(__data, *args, **kwargs):
 
 # option: no args, custom dispatch (e.g. register NoArgs)
 # strips symbols
+@overload
+def verb_dispatch(cls: "_DispatchType") -> "_DispatchDecorator":
+    ...
+@overload
+def verb_dispatch(cls: "_DispatchType", f: Callable) -> "_SingleDispatchCallable[Any]":
+    ...
 def verb_dispatch(cls, f = None):
     """Wrap singledispatch. Making sure to keep its attributes on the wrapper.
 
@@ -141,7 +154,7 @@ def verb_dispatch(cls, f = None):
 
     # classic way of allowing args to a decorator
     if f is None:
-        return lambda f: verb_dispatch(cls, f)
+        return partial(verb_dispatch, cls)
 
     # initially registers func for object, so need to change to pd.DataFrame
     dispatch_func = singledispatch(f)