cleanup

Author: seeM

fastcore/transform.py

@@ -15,60 +15,51 @@
 from .dispatch import *
 import inspect
 from copy import copy
-from plum import add_conversion_method, dispatch, Dispatcher
-
-# Cell
-# TODO: Shouldn't sig of method set first parameter type to self? how do i get that...
-#       i.e. self assumes type(self) / self.__class__?
-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])
-    # TODO: Shouldn't we create a Function here and dispatch from that?
-    #       We don't need this in the dispatch table i think...
-    # TODO: Should Function take name qualname etc and have a .from_callable method?
-    #       since the func isn't even dispatched by default
-    pf = d(f)
-    if cls is not None:
-        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
+from plum import add_conversion_method, dispatch, Function
 
 # Cell
 _tfm_methods = 'encodes','decodes','setups'
 
 def _is_tfm_method(f,n):
     return n in _tfm_methods and callable(f)
 
-# TODO: Do we still need this given the fact that plum searches mro without needing them to be Functions?
 class _TfmDict(dict):
     def __setitem__(self,k,v):
         if _is_tfm_method(v,k): v = dispatch(v)
         super().__setitem__(k,v)
 
+# Cell
+def _dispatch(f, cls):
+    "Dispatch and set a function as an instance method"
+    f = copy(f)
+    n = f.__name__
+    # plum uses __qualname__ to infer f's owner
+    f.__qualname__ = f'{cls.__name__}.{n}'
+    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
+
 # Cell
 class _TfmMeta(type):
-    # TODO: commenting since this breaks inspect.signature of Transform instances,
-    #       which then breaks inspect.signature of a partial of a Transform instance
-    #def __new__(cls, name, bases, dict):
-    #    res = super().__new__(cls, name, bases, dict)
-    #    res.__signature__ = inspect.signature(res.__init__)
-    #    return res
-
-    # TODO: Can we move this to Transform.__init__? Then we don't reeeeeally need a metaclass anymore...?
-    #       Ohhhhh man, can we dispatch this? If called with callable, do this, and so on
+    def __new__(cls, name, bases, dict):
+        # _TfmMeta.__call__ shadows the signature of inheriting classes, set it back
+        res = super().__new__(cls, name, bases, dict)
+        res.__signature__ = inspect.signature(res.__init__)
+        return res
+
     def __call__(cls, *args, **kwargs):
         f = first(args)
-        n = getattr(f,'__name__',None)
-        if _is_tfm_method(f,n):
-            # use __dict__ over hasattr since it excludes parent classes
-            if n in cls.__dict__: return getattr(cls,n).dispatch(f)
-            return _mk_plum_func(dispatch,n,f,cls)
-        return super().__call__(*args,**kwargs)
+        n = getattr(f, '__name__', None)
+        if _is_tfm_method(f, n): return _dispatch(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__)
+        return obj
 
     @classmethod
     def __prepare__(cls, name, bases): return _TfmDict()
@@ -82,6 +73,21 @@ def _get_name(o):
 # Cell
 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())
+
+# Cell
+def _pf_repr(pf): return '\n'.join(f"{f.__name__}: ({','.join(_pt_repr(t) for t in s.types)}) -> {_pt_repr(r)}"
+                                   for s, (f, r) in pf.methods.items())
+
 # Cell
 class Transform(metaclass=_TfmMeta):
     "Delegates (`__call__`,`decode`,`setup`) to (<code>encodes</code>,<code>decodes</code>,<code>setups</code>) if `split_idx` matches"
@@ -92,10 +98,8 @@ 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._d = Dispatcher() # TODO: do we need to hold this reference?
-        # TODO: do u have to set the name from the original func here? for pipelines to work
-        # TODO: I don't think this is registering any methods! Why is a func needed then???
-        for n in _tfm_methods: setattr(self,n,_mk_plum_func(self._d,n))
+        def identity(x): return x
+        for n in _tfm_methods: setattr(self,n,Function(identity).dispatch(identity))
         if enc:
             self.encodes.dispatch(enc)
             self.order = getattr(enc,'order',self.order)
@@ -107,18 +111,7 @@ def __init__(self, enc=None, dec=None, split_idx=None, order=None):
     def name(self): return getattr(self, '_name', _get_name(self))
     def __call__(self, x, **kwargs): return self._call('encodes', x, **kwargs)
     def decode  (self, x, **kwargs): return self._call('decodes', x, **kwargs)
-    def __repr__(self):
-        def _pf_repr(pf):
-            e = []
-            for s, (f_, r) in self.encodes.methods.items():
-                types = ','.join(str(o) for o in s.types)
-                types = f'({types})'
-                e.append(f'{f_.__name__}: {types} -> {r}')
-            return e
-        r = f'{self.name}:'
-        r += '\n  encodes:\n' + '\n'.join('    ' + o for o in _pf_repr(self.encodes))
-        r += '\n  decodes:\n' + '\n'.join('    ' + o for o in _pf_repr(self.decodes))
-        return r
+    def __repr__(self): return f'{self.name}:\nencodes: {_pf_repr(self.encodes)}\ndecodes: {_pf_repr(self.decodes)}'
 
     def setup(self, items=None, train_setup=False):
         train_setup = train_setup if self.train_setup is None else self.train_setup
@@ -134,7 +127,7 @@ def _do_call(self, f, x, **kwargs):
             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
+            # plum reads empty return annotation as object, retain_type expects it 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)
@@ -146,7 +139,7 @@ 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
+#Implement the Transform convention that a None return annotation disables conversion
 add_conversion_method(object, NoneType, lambda x: x)
 
 # Cell