core

plum-dispatch extensions

fastdispatch extends the wonderful plum’s Julia-inspired implementation of multiple dispatch for Python.

import numpy as np
from fastcore.test import *

---

FastFunction

 FastFunction (f, owner=None)

Multiple dispatched function; extends plum.Function

It has a concise repr:

def f(x:int) -> float: pass
f = FastFunction(f).dispatch(f)
f

f(int) -> float

It supports fastcore’s backport of the | operator on types:

def f1(x):          return 'obj'
def f2(x:int|str): return 'int|str'
f = FastFunction(f1).dispatch(f1).dispatch(f2)

test_eq(f(0),   'int|str')
test_eq(f(''),  'int|str')
test_eq(f(0.0), 'obj')

Indexing a FastFunction works like plum.Function.invoke but returns the most-specific matching method with the fewest parameters:

def f1(a:int,   b,     c):    return 'int, 3 args'
def f2(a:int,   b,     c, d): return 'int, 4 args'
def f3(a:float, b,     c):    return 'float, 3 args'
def f4(a:float, b:str, c):    return 'float, str, 3 args'
f = FastFunction(f1).dispatch(f1).dispatch(f2).dispatch(f3).dispatch(f4)

test_eq(f[int](0,0,0),        'int, 3 args')
test_eq(f[float](0,0,0),      'float, 3 args')
test_eq(f[float](0,0,0),      'float, 3 args')
test_eq(f[float, str](0,0,0), 'float, str, 3 args')

---

FastDispatcher

 FastDispatcher ()

Namespace for multiple dispatched functions; extends plum.Dispatcher

dispatch = FastDispatcher()

Dispatching with FastDispatcher returns a FastFunction:

@dispatch
def f(x): return 'obj'

assert isinstance(f, FastFunction)

It supports fastcore’s backport of the | operator on types:

@dispatch
def f(x:int|str): return 'int|str'

test_eq(f(0),   'int|str')
test_eq(f(''),  'int|str')
test_eq(f(0.0), 'obj')

… FastDispatcher.multi works too:

@dispatch.multi([bool],[list])
def f(x:bool|list): return 'bool|list'
@dispatch
def f(x:int): return 'int'

test_eq(f(True), 'bool|list')
test_eq(f([]),   'bool|list')
test_eq(f(0),    'int')

---

FastDispatcher.to

 FastDispatcher.to (cls)

Decorator: dispatch f to cls.f

This lets you dynamically extend dispatched methods:

class A:
    @dispatch
    def f(self, x): return 'obj'

@dispatch.to(A)
def f(self, x:int): return 'int'

a = A()
test_eq(a.f(0), 'int')
test_eq(a.f(''), 'obj')

Casting

Now that we can dispatch on types, let’s make it easier to cast objects to a different type.

---

retain_meta

 retain_meta (x, res, as_copy=False)

Call res.set_meta(x), if it exists

---

cast

 cast (x, typ)

Cast x to typ (may change x inplace)

This works both for plain python classes:…

mk_class('_T1', 'a')      # mk_class is a fastcore utility that constructs a class
class _T2(_T1): pass

t = _T1(a=1)
t2 = cast(t, _T2)        
assert t2 is t            # t2 refers to the same object as t
assert isinstance(t,_T2)  # t also changed in-place
assert isinstance(t2,_T2)

test_eq_type(_T2(a=1), t2)

…as well as for arrays and tensors.

class _T1(np.ndarray): pass

t = np.array([1])
t2 = cast(t, _T1)
test_eq(np.array([1]), t2)
test_eq(_T1, type(t2))

To customize casting for other types, define a separate cast function with dispatch for your type.

---

retain_type

 retain_type (new, old=None, typ=None, as_copy=False)

Cast new to type of old or typ if it’s a superclass

class _T(tuple): pass
a = _T((1,2))
b = tuple((1,2))
c = retain_type(b, typ=_T)
test_eq_type(c, a)

If old has a _meta attribute, its content is passed when casting new to the type of old. In the below example, only the attribute a, but not other_attr is kept, because other_attr is not in _meta:

def default_set_meta(self, x, as_copy=False):
    "Copy over `_meta` from `x` to `res`, if it's missing"
    if hasattr(x, '_meta') and not hasattr(self, '_meta'):
        meta = x._meta
        if as_copy: meta = copy(meta)
        self._meta = meta
    return self

class _A():
    set_meta = default_set_meta
    def __init__(self, t): self.t=t

class _B1(_A):
    def __init__(self, t, a=1):
        super().__init__(t)
        self._meta = {'a':a}
        self.other_attr = 'Hello' # will not be kept after casting.
        
x = _B1(1, a=2)
b = _A(1)
c = retain_type(b, old=x)
test_eq(c._meta, {'a': 2})
assert not getattr(c, 'other_attr', None)

---

retain_types

 retain_types (new, old=None, typs=None)

Cast each item of new to type of matching item in old if it’s a superclass

class T(tuple): pass

t1,t2 = retain_types((1,(1,(1,1))), (2,T((2,T((3,4))))))
test_eq_type(t1, 1)
test_eq_type(t2, T((1,T((1,1)))))

t1,t2 = retain_types((1,(1,(1,1))), typs = {tuple: [int, {T: [int, {T: [int,int]}]}]})
test_eq_type(t1, 1)
test_eq_type(t2, T((1,T((1,1)))))

---

explode_types

 explode_types (o)

Return the type of o, potentially in nested dictionaries for thing that are listy

test_eq(explode_types((2,T((2,T((3,4)))))), {tuple: [int, {T: [int, {T: [int,int]}]}]})