works with tupple annotations; working autoreload; repr

Author: seeM

fastcore/imports.py

@@ -14,6 +14,15 @@
     MethodDescriptorType = type(str.join)
 from types import BuiltinFunctionType,BuiltinMethodType,MethodType,FunctionType,SimpleNamespace
 
+#Patch autoreload (if its loaded) to work with plum
+try: from IPython import get_ipython
+except ImportError: pass
+else:
+    ip = get_ipython()
+    if ip is not None and 'IPython.extensions.storemagic' in ip.extension_manager.loaded:
+        from plum.autoreload import activate
+        activate()
+
 NoneType = type(None)
 string_classes = (str,bytes)
 

fastcore/transform.py

@@ -16,33 +16,42 @@
 import inspect
 from copy import copy
 from plum import add_conversion_method, dispatch, Function
+from typing import get_args, get_origin
 
 # Cell
-_tfm_methods = 'encodes','decodes','setups'
-
-def _is_tfm_method(f,n):
-    return n in _tfm_methods and callable(f)
+# Convert tuple annotations to unions to work with plum
+def _annot_tuple_to_union(f):
+    for k, v in type_hints(f).items():
+        if isinstance(v, tuple): f.__annotations__[k] = Union[v]
+    return f
 
-class _TfmDict(dict):
-    def __setitem__(self,k,v):
-        if _is_tfm_method(v,k): v = dispatch(v)
-        super().__setitem__(k,v)
+def _dispatch(f): return dispatch(_annot_tuple_to_union(f))
 
-# Cell
-def _dispatch(f, cls):
-    "Dispatch and set a function as an instance method"
+def _dispatch_method(f, cls):
     f = copy(f)
     n = f.__name__
     # plum uses __qualname__ to infer f's owner
     f.__qualname__ = f'{cls.__name__}.{n}'
-    pf = dispatch(f)
+    pf = _dispatch(f)
     setattr(cls, n, pf)
     # plum uses __set_name__ to resolve a plum.Function's owner.
     # since we assign after class creation, __set_name__ must be called directly
     # source: https://docs.python.org/3/reference/datamodel.html#object.__set_name__
     pf.__set_name__(cls, n)
     return pf
 
+def _pf_dispatch(pf, f): return pf.dispatch(_annot_tuple_to_union(f))
+
+# Cell
+_tfm_methods = 'encodes','decodes','setups'
+
+def _is_tfm_method(f, n): return n in _tfm_methods and callable(f)
+
+class _TfmDict(dict):
+    def __setitem__(self, k, v):
+        if _is_tfm_method(v, k): v = _dispatch(v)
+        super().__setitem__(k, v)
+
 # Cell
 class _TfmMeta(type):
     def __new__(cls, name, bases, dict):
@@ -54,8 +63,8 @@ def __new__(cls, name, bases, dict):
     def __call__(cls, *args, **kwargs):
         f = first(args)
         n = getattr(f, '__name__', None)
-        if _is_tfm_method(f, n): return _dispatch(f, cls)
-        obj = super().__call__(*args,**kwargs)
+        if _is_tfm_method(f, n): return _dispatch_method(f, cls)
+        obj = super().__call__(*args, **kwargs)
         # _TfmMeta.__new__ replaces cls.__signature__ which breaks the signature of a callable
         # instances of cls, fix it
         if hasattr(obj, '__call__'): obj.__signature__ = inspect.signature(obj.__call__)
@@ -74,15 +83,16 @@ def _get_name(o):
 def _is_tuple(o): return isinstance(o, tuple) and not hasattr(o, '_fields')
 
 # Cell
-@dispatch
 def _pt_repr(o):
     n = type(o).__name__
     if n == 'Tuple': return f"{n.lower()}[{','.join(_pt_repr(t) for t in o._el_types)}]"
     if n == 'List': return f'{n.lower()}[{_pt_repr(o._el_type)}]'
     if n == 'Dict': return f'{n.lower()}[{_pt_repr(o._key_type)},{_pt_repr(o._value_type)}]'
     if n in ('Sequence','Iterable'): return f'{n}[{_pt_repr(o._el_type)}]'
-    if n == 'VarArgs': return f'VarArgs[{_pt_repr(o.type)}]'
-    return '|'.join(t.__name__ for t in o.get_types())
+    if n == 'VarArgs': return f'{n}[{_pt_repr(o.type)}]'
+    if n == 'Union': return '|'.join(sorted(t.__name__ for t in (o.get_types())))
+    assert len(o.get_types()) == 1
+    return o.get_types()[0].__name__
 
 # Cell
 def _pf_repr(pf): return '\n'.join(f"{f.__name__}: ({','.join(_pt_repr(t) for t in s.types)}) -> {_pt_repr(r)}"
@@ -101,11 +111,14 @@ def __init__(self, enc=None, dec=None, split_idx=None, order=None):
         def identity(x): return x
         for n in _tfm_methods: setattr(self,n,Function(identity).dispatch(identity))
         if enc:
-            self.encodes.dispatch(enc)
+            _pf_dispatch(self.encodes, enc)
             self.order = getattr(enc,'order',self.order)
-            if len(type_hints(enc)) > 0: self.input_types = first(type_hints(enc).values())
+            if len(type_hints(enc)) > 0:
+                self.input_types = first(type_hints(enc).values())
+                # Convert Union to tuple, remove once the rest of fastai supports Union
+                if get_origin(self.input_types) is Union: self.input_types=get_args(self.input_types)
             self._name = _get_name(enc)
-        if dec: self.decodes.dispatch(dec)
+        if dec: _pf_dispatch(self.decodes, dec)
 
     @property
     def name(self): return getattr(self, '_name', _get_name(self))

nbs/05_transform.ipynb

@@ -23,7 +23,8 @@
     "from fastcore.dispatch import *\n",
     "import inspect\n",
     "from copy import copy\n",
-    "from plum import add_conversion_method, dispatch, Function"
+    "from plum import add_conversion_method, dispatch, Function\n",
+    "from typing import get_args, get_origin"
    ]
   },
   {
@@ -69,37 +70,45 @@
    "outputs": [],
    "source": [
     "#export\n",
-    "_tfm_methods = 'encodes','decodes','setups'\n",
+    "# Convert tuple annotations to unions to work with plum\n",
+    "def _annot_tuple_to_union(f):\n",
+    "    for k, v in type_hints(f).items():\n",
+    "        if isinstance(v, tuple): f.__annotations__[k] = Union[v]\n",
+    "    return f\n",
     "\n",
-    "def _is_tfm_method(f,n):\n",
-    "    return n in _tfm_methods and callable(f)\n",
+    "def _dispatch(f): return dispatch(_annot_tuple_to_union(f))\n",
     "\n",
-    "class _TfmDict(dict):\n",
-    "    def __setitem__(self,k,v):\n",
-    "        if _is_tfm_method(v,k): v = dispatch(v)\n",
-    "        super().__setitem__(k,v)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "#export\n",
-    "def _dispatch(f, cls):\n",
-    "    \"Dispatch and set a function as an instance method\"\n",
+    "def _dispatch_method(f, cls):\n",
     "    f = copy(f)\n",
     "    n = f.__name__\n",
     "    # plum uses __qualname__ to infer f's owner\n",
     "    f.__qualname__ = f'{cls.__name__}.{n}'\n",
-    "    pf = dispatch(f)\n",
+    "    pf = _dispatch(f)\n",
     "    setattr(cls, n, pf)\n",
     "    # plum uses __set_name__ to resolve a plum.Function's owner.\n",
     "    # since we assign after class creation, __set_name__ must be called directly\n",
     "    # source: https://docs.python.org/3/reference/datamodel.html#object.__set_name__\n",
     "    pf.__set_name__(cls, n)\n",
-    "    return pf"
+    "    return pf\n",
+    "\n",
+    "def _pf_dispatch(pf, f): return pf.dispatch(_annot_tuple_to_union(f))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#export\n",
+    "_tfm_methods = 'encodes','decodes','setups'\n",
+    "\n",
+    "def _is_tfm_method(f, n): return n in _tfm_methods and callable(f)\n",
+    "\n",
+    "class _TfmDict(dict):\n",
+    "    def __setitem__(self, k, v):\n",
+    "        if _is_tfm_method(v, k): v = _dispatch(v)\n",
+    "        super().__setitem__(k, v)"
    ]
   },
   {
@@ -119,8 +128,8 @@
     "    def __call__(cls, *args, **kwargs):\n",
     "        f = first(args)\n",
     "        n = getattr(f, '__name__', None)\n",
-    "        if _is_tfm_method(f, n): return _dispatch(f, cls)\n",
-    "        obj = super().__call__(*args,**kwargs)\n",
+    "        if _is_tfm_method(f, n): return _dispatch_method(f, cls)\n",
+    "        obj = super().__call__(*args, **kwargs)\n",
     "        # _TfmMeta.__new__ replaces cls.__signature__ which breaks the signature of a callable\n",
     "        # instances of cls, fix it\n",
     "        if hasattr(obj, '__call__'): obj.__signature__ = inspect.signature(obj.__call__)\n",
@@ -160,15 +169,16 @@
    "outputs": [],
    "source": [
     "#export\n",
-    "@dispatch\n",
     "def _pt_repr(o):\n",
     "    n = type(o).__name__\n",
     "    if n == 'Tuple': return f\"{n.lower()}[{','.join(_pt_repr(t) for t in o._el_types)}]\"\n",
     "    if n == 'List': return f'{n.lower()}[{_pt_repr(o._el_type)}]'\n",
     "    if n == 'Dict': return f'{n.lower()}[{_pt_repr(o._key_type)},{_pt_repr(o._value_type)}]'\n",
     "    if n in ('Sequence','Iterable'): return f'{n}[{_pt_repr(o._el_type)}]'\n",
-    "    if n == 'VarArgs': return f'VarArgs[{_pt_repr(o.type)}]'\n",
-    "    return '|'.join(t.__name__ for t in o.get_types())"
+    "    if n == 'VarArgs': return f'{n}[{_pt_repr(o.type)}]'\n",
+    "    if n == 'Union': return '|'.join(sorted(t.__name__ for t in (o.get_types())))\n",
+    "    assert len(o.get_types()) == 1\n",
+    "    return o.get_types()[0].__name__"
    ]
   },
   {
@@ -188,7 +198,8 @@
     "test_eq(_pt_repr(ptype(Iterable[int])), 'Iterable[int]')\n",
     "test_eq(_pt_repr(ptype(Dict[str, int])), 'dict[str,int]')\n",
     "test_eq(_pt_repr(ptype(VarArgs[str])), 'VarArgs[str]')\n",
-    "test_eq(_pt_repr(ptype(Dict[Tuple[Union[int,str],float], List[Tuple[object]]])), 'dict[tuple[int|str,float],list[tuple[object]]]')"
+    "test_eq(_pt_repr(ptype(Dict[Tuple[Union[int,str],float], List[Tuple[object]]])),\n",
+    "        'dict[tuple[int|str,float],list[tuple[object]]]')"
    ]
   },
   {
@@ -233,11 +244,14 @@
     "        def identity(x): return x\n",
     "        for n in _tfm_methods: setattr(self,n,Function(identity).dispatch(identity))\n",
     "        if enc:\n",
-    "            self.encodes.dispatch(enc)\n",
+    "            _pf_dispatch(self.encodes, enc)\n",
     "            self.order = getattr(enc,'order',self.order)\n",
-    "            if len(type_hints(enc)) > 0: self.input_types = first(type_hints(enc).values())\n",
+    "            if len(type_hints(enc)) > 0:\n",
+    "                self.input_types = first(type_hints(enc).values())\n",
+    "                # Convert Union to tuple, remove once the rest of fastai supports Union\n",
+    "                if get_origin(self.input_types) is Union: self.input_types=get_args(self.input_types)\n",
     "            self._name = _get_name(enc)\n",
-    "        if dec: self.decodes.dispatch(dec)\n",
+    "        if dec: _pf_dispatch(self.decodes, dec)\n",
     "\n",
     "    @property\n",
     "    def name(self): return getattr(self, '_name', _get_name(self))\n",
@@ -292,7 +306,7 @@
       "text/markdown": [
        "<h2 id=\"Transform\" class=\"doc_header\"><code>class</code> <code>Transform</code><a href=\"\" class=\"source_link\" style=\"float:right\">[source]</a></h2>\n",
        "\n",
-       "> <code>Transform</code>(**`enc`**=*`None`*, **`dec`**=*`None`*, **`split_idx`**=*`None`*, **`order`**=*`None`*)\n",
+       "> <code>Transform</code>(**`self`**, **`enc`**=*`None`*, **`dec`**=*`None`*, **`split_idx`**=*`None`*, **`order`**=*`None`*)\n",
        "\n",
        "Delegates (`__call__`,`decode`,`setup`) to (<code>encodes</code>,<code>decodes</code>,<code>setups</code>) if `split_idx` matches"
       ],
@@ -473,6 +487,25 @@
     "test_eq_type(f3(2), 2)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Transforms can be created from class methods too:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class A:\n",
+    "    @classmethod\n",
+    "    def create(cls, x:int): return x+1\n",
+    "test_eq(Transform(A.create)(1), 2)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -607,8 +640,8 @@
     "class MyClass(int): pass\n",
     "\n",
     "class A(Transform):\n",
-    "    def encodes(self, x:Union[MyClass,float]): return x/2\n",
-    "    def encodes(self, x:Union[str,list]): return str(x)+'_1'\n",
+    "    def encodes(self, x:(MyClass,float)): return x/2\n",
+    "    def encodes(self, x:(str,list)): return str(x)+'_1'\n",
     "\n",
     "f = A()"
    ]
@@ -647,6 +680,18 @@
     "test_eq(f(['a','b','c']), \"['a', 'b', 'c']_1\") # input is of type list"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "@Transform\n",
+    "def f(x:(int,float)): return x+1\n",
+    "test_eq(f(0), 1)\n",
+    "test_eq(f('a'), 'a')"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -901,7 +946,7 @@
     {
      "data": {
       "text/plain": [
-       "Promise(obj=<function <function AL.encodes at 0x1220569d0> with 2 method(s)>)"
+       "Promise(obj=<function <function AL.encodes at 0x11e3fb670> with 2 method(s)>)"
       ]
      },
      "execution_count": null,
@@ -1364,20 +1409,6 @@
     "test_eq(type(f.decode((1,1))), _T)"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "#### Transform tests -"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
   {
    "cell_type": "markdown",
    "metadata": {},