prototype plum-dispatch for fastcore.transform

Author: seeM

fastcore/transform.py

@@ -9,37 +9,50 @@
 from .utils import *
 from .dispatch import *
 import inspect
+from plum import add_conversion_method, dispatch, Dispatcher
+
+# Cell
+def _mk_plum_func(d, n, f=None, cls=None):
+    f = (lambda x: x) if f is None else copy(f)
+    f.__name__ = n
+    # plum uses __qualname__ to infer f's owner
+    f.__qualname__ = n if cls is None else '.'.join([cls.__name__,n])
+    pf = d(f)
+    if cls is not None:
+        setattr(cls,n,pf)
+        # plum uses __set_name__ to resolve a Function's owner.
+        # since class variable is assigned after class is created, __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
 
 # 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 k not in _tfm_methods or not callable(v): return super().__setitem__(k,v)
-        if k not in self: super().__setitem__(k,TypeDispatch())
-        self[k].add(v)
+        if _is_tfm_method(v,k): v = dispatch(v)
+        super().__setitem__(k,v)
 
 # Cell
 class _TfmMeta(type):
+    @classmethod
+    def __prepare__(cls, name, bases): return _TfmDict()
+
     def __new__(cls, name, bases, dict):
         res = super().__new__(cls, name, bases, dict)
-        for nm in _tfm_methods:
-            base_td = [getattr(b,nm,None) for b in bases]
-            if nm in res.__dict__: getattr(res,nm).bases = base_td
-            else: setattr(res, nm, TypeDispatch(bases=base_td))
         res.__signature__ = inspect.signature(res.__init__)
         return res
 
     def __call__(cls, *args, **kwargs):
-        f = args[0] if args else None
+        f = first(args)
         n = getattr(f,'__name__',None)
-        if callable(f) and n in _tfm_methods:
-            getattr(cls,n).add(f)
-            return f
-        return super().__call__(*args, **kwargs)
-
-    @classmethod
-    def __prepare__(cls, name, bases): return _TfmDict()
+        if not _is_tfm_method(f,n): return super().__call__(*args,**kwargs)
+        if n in cls.__dict__: getattr(cls,n).dispatch(f)
+        return _mk_plum_func(dispatch,n,f,cls=cls)
 
 # Cell
 def _get_name(o):
@@ -60,13 +73,14 @@ def __init__(self, enc=None, dec=None, split_idx=None, order=None):
         self.init_enc = enc or dec
         if not self.init_enc: return
 
-        self.encodes,self.decodes,self.setups = TypeDispatch(),TypeDispatch(),TypeDispatch()
+        self._d = Dispatcher() # TODO: do we need to hold this reference?
+        for n in _tfm_methods: setattr(self,n,_mk_plum_func(self._d,n))
         if enc:
-            self.encodes.add(enc)
+            self.encodes = self.encodes.dispatch(enc)
             self.order = getattr(enc,'order',self.order)
             if len(type_hints(enc)) > 0: self.input_types = first(type_hints(enc).values())
             self._name = _get_name(enc)
-        if dec: self.decodes.add(dec)
+        if dec: self.decodes = self.decodes.dispatch(dec)
 
     @property
     def name(self): return getattr(self, '_name', _get_name(self))
@@ -85,13 +99,24 @@ def _call(self, fn, x, split_idx=None, **kwargs):
     def _do_call(self, f, x, **kwargs):
         if not _is_tuple(x):
             if f is None: return x
-            ret = f.returns(x) if hasattr(f,'returns') else None
-            return retain_type(f(x, **kwargs), x, ret)
+            ts = [type(self),type(x)] if hasattr(f,'instance') else [type(x)]
+            _, ret = f.resolve_method(*ts)
+            ret = ret._type
+            # plum reads empty return annot as object, fastcore reads as None
+            if ret is object: ret = None
+            return retain_type(f(x,**kwargs), x, ret)
         res = tuple(self._do_call(f, x_, **kwargs) for x_ in x)
         return retain_type(res, x)
+    def encodes(self, x): return x
+    def decodes(self, x): return x
+    def setups(self, dl): return dl
 
 add_docs(Transform, decode="Delegate to <code>decodes</code> to undo transform", setup="Delegate to <code>setups</code> to set up transform")
 
+# Cell
+#Transform interpret's None return type as no conversion
+add_conversion_method(object, NoneType, lambda x: x)
+
 # Cell
 class InplaceTransform(Transform):
     "A `Transform` that modifies in-place and just returns whatever it's passed"

nbs/05_transform.ipynb

@@ -26,7 +26,7 @@
 min_python = cfg['min_python']
 lic = licenses[cfg['license']]
 
-requirements = ['pip', 'packaging']
+requirements = ['pip', 'packaging', 'plum-dispatch']
 if cfg.get('requirements'): requirements += cfg.get('requirements','').split()
 if cfg.get('pip_requirements'): requirements += cfg.get('pip_requirements','').split()
 dev_requirements = (cfg.get('dev_requirements') or '').split()