multimethod.py

from collections import defaultdict
import functools


class Registry(object):
    def __init__(self):
        self.dispatchers = {}
        self.functions = defaultdict(dict)

    def __enter__(self):
        self.old_registry = multimethod.registry
        multimethod.registry = self

    def __exit__(self, *args):
        multimethod.registry = self.old_registry


class multimethod(object):

    registry = Registry()

    def __init__(self, dispatch_value):
        self._dispatch_value = dispatch_value

    def __call__(self, function):
        registry = self.registry
        registry.functions[function.__name__][self._dispatch_value] = function
        def proxy(*args, **kwargs):
            dispatch_value = registry.dispatchers[function.__name__](*args, **kwargs)
            return registry.functions[function.__name__][dispatch_value](*args, **kwargs)
        return proxy

    @classmethod
    def register(cls, name, dispatcher=None):
        if dispatcher is not None:
            cls.registry.dispatchers[name] = dispatcher
        else:
            return functools.partial(cls.register, name)

"""
>>> with Registry():
...     @multimethod.register('add_event')
...     def dispatch(event):
...         return event['type']
...
...     @multimethod('game')
...     def add_event(event):
...         print 'GAME'
...
...     @multimethod('event')
...     def add_event(event):
...         print 'EVENT'
...
...     add_event({'type': 'game'})
GAME
>>> add_event({'type': 'game'})
GAME
"""