Class Dispatch

dask/utils.py:709–789

Simple single dispatch.

class Dispatch:
    """Simple single dispatch."""

    def __init__(self, name=None):
        self._lookup = {}
        self._lazy = {}
        if name:
            self.__name__ = name

    def register(self, type, func=None):
        """Register dispatch of `func` on arguments of type `type`"""

        def wrapper(func):
            if isinstance(type, tuple):
                for t in type:
                    self.register(t, func)
            else:
                self._lookup[type] = func
            return func

        return wrapper(func) if func is not None else wrapper

    def register_lazy(self, toplevel, func=None):
        """
        Register a registration function which will be called if the
        *toplevel* module (e.g. 'pandas') is ever loaded.
        """

        def wrapper(func):
            self._lazy[toplevel] = func
            return func

        return wrapper(func) if func is not None else wrapper

    def dispatch(self, cls):
        """Return the function implementation for the given ``cls``"""
        lk = self._lookup
        if cls in lk:
            return lk[cls]
        for cls2 in cls.__mro__:
            # Is a lazy registration function present?
            try:
                toplevel, _, _ = cls2.__module__.partition(".")
            except Exception:
                continue
            try:
                register = self._lazy[toplevel]
            except KeyError:
                pass
            else:
                register()
                self._lazy.pop(toplevel, None)
                meth = self.dispatch(cls)  # recurse
                lk[cls] = meth
                lk[cls2] = meth
                return meth
            try:
                impl = lk[cls2]