Merge pull request #106 from Salehbigdeli/master

Make typedispatch decorator more general

Author: jph00

fastcore/dispatch.py

@@ -97,7 +97,8 @@ def __init__(self, funcs=(), bases=()):
 
     def add(self, f):
         "Add type `t` and function `f`"
-        a0,a1 = _p2_anno(f)
+        if isinstance(f, staticmethod): a0,a1 = _p2_anno(f.__func__)
+        else: a0,a1 = _p2_anno(f)
         t = self.funcs.d.get(a0)
         if t is None:
             t = _TypeDict()
@@ -121,7 +122,8 @@ def __call__(self, *args, **kwargs):
         ts = L(args).map(type)[:2]
         f = self[tuple(ts)]
         if not f: return args[0]
-        if self.inst is not None: f = MethodType(f, self.inst)
+        if isinstance(f, staticmethod): f = f.__func__
+        elif self.inst is not None: f = MethodType(f, self.inst)
         elif self.owner is not None: f = MethodType(f, self.owner)
         return f(*args, **kwargs)
 
@@ -152,7 +154,9 @@ class DispatchReg:
     "A global registry for `TypeDispatch` objects keyed by function name"
     def __init__(self): self.d = defaultdict(TypeDispatch)
     def __call__(self, f):
-        nm = f'{f.__qualname__}'
+        if isinstance(f, (classmethod, staticmethod)): nm = f'{f.__func__.__qualname__}'
+        else: nm = f'{f.__qualname__}'
+        if isinstance(f, classmethod): f=f.__func__
         self.d[nm].add(f)
         return self.d[nm]
 

nbs/03_dispatch.ipynb

@@ -449,7 +449,8 @@
     "\n",
     "    def add(self, f):\n",
     "        \"Add type `t` and function `f`\"\n",
-    "        a0,a1 = _p2_anno(f)\n",
+    "        if isinstance(f, staticmethod): a0,a1 = _p2_anno(f.__func__)\n",
+    "        else: a0,a1 = _p2_anno(f)\n",
     "        t = self.funcs.d.get(a0)\n",
     "        if t is None:\n",
     "            t = _TypeDict()\n",
@@ -473,7 +474,8 @@
     "        ts = L(args).map(type)[:2]\n",
     "        f = self[tuple(ts)]\n",
     "        if not f: return args[0]\n",
-    "        if self.inst is not None: f = MethodType(f, self.inst)\n",
+    "        if isinstance(f, staticmethod): f = f.__func__\n",
+    "        elif self.inst is not None: f = MethodType(f, self.inst)\n",
     "        elif self.owner is not None: f = MethodType(f, self.owner)\n",
     "        return f(*args, **kwargs)\n",
     "\n",
@@ -943,7 +945,7 @@
     {
      "data": {
       "text/markdown": [
-       "<h4 id=\"TypeDispatch.__call__\" class=\"doc_header\"><code>TypeDispatch.__call__</code><a href=\"__main__.py#L33\" class=\"source_link\" style=\"float:right\">[source]</a></h4>\n",
+       "<h4 id=\"TypeDispatch.__call__\" class=\"doc_header\"><code>TypeDispatch.__call__</code><a href=\"__main__.py#L34\" class=\"source_link\" style=\"float:right\">[source]</a></h4>\n",
        "\n",
        "> <code>TypeDispatch.__call__</code>(**\\*`args`**, **\\*\\*`kwargs`**)\n",
        "\n",
@@ -1031,7 +1033,7 @@
     {
      "data": {
       "text/markdown": [
-       "<h4 id=\"TypeDispatch.returns\" class=\"doc_header\"><code>TypeDispatch.returns</code><a href=\"__main__.py#L21\" class=\"source_link\" style=\"float:right\">[source]</a></h4>\n",
+       "<h4 id=\"TypeDispatch.returns\" class=\"doc_header\"><code>TypeDispatch.returns</code><a href=\"__main__.py#L22\" class=\"source_link\" style=\"float:right\">[source]</a></h4>\n",
        "\n",
        "> <code>TypeDispatch.returns</code>(**`x`**)\n",
        "\n",
@@ -1198,7 +1200,9 @@
     "    \"A global registry for `TypeDispatch` objects keyed by function name\"\n",
     "    def __init__(self): self.d = defaultdict(TypeDispatch)\n",
     "    def __call__(self, f):\n",
-    "        nm = f'{f.__qualname__}'\n",
+    "        if isinstance(f, (classmethod, staticmethod)): nm = f'{f.__func__.__qualname__}'\n",
+    "        else: nm = f'{f.__qualname__}'\n",
+    "        if isinstance(f, classmethod): f=f.__func__\n",
     "        self.d[nm].add(f)\n",
     "        return self.d[nm]\n",
     "\n",
@@ -1227,6 +1231,36 @@
     "test_eq(f_td_test('a','b'), 'ab')"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### Using typedispatch With other decorators\n",
+    "\n",
+    "You can use `typedispatch` with `classmethod` and `staticmethod` decorator"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class A:\n",
+    "    @typedispatch\n",
+    "    def f_td_test(self, x:numbers.Integral, y): return x+1\n",
+    "    @typedispatch\n",
+    "    @classmethod\n",
+    "    def f_td_test(cls, x:int, y:float): return x+y\n",
+    "    @typedispatch\n",
+    "    @staticmethod\n",
+    "    def f_td_test(x:int, y:int): return x*y\n",
+    "    \n",
+    "test_eq(A.f_td_test(3,2), 6)\n",
+    "test_eq(A.f_td_test(3,2.0), 5)\n",
+    "test_eq(A().f_td_test(3,'2.0'), 4)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},