Multiple dispatch decorator in Python

Question

I've written a decorator for easily creating multiple dispatch functions:

from functools import wraps
def multi_dispatch(for_function):
 """
 Returns a multiple dispatch version of the function.
 The returned function doesn't allow keyword arguments.
 >>> @multi_dispatch
 ... def foo(a, b):
 ... pass
 ...
 >>> @foo.add
 ... def _foo(a: int, b: int, c : str = '33'):
 ... return a + b*3 + len(c)
 ...
 >>> @foo.add
 ... def _foo(a: int, b: int, c = 3):
 ... return a + b*3 - c
 ...
 >>> @foo.add
 ... def _foo(a: float, b: float):
 ... return a*2 - b
 ...
 >>> foo(4, 5)
 16
 >>> foo(1.0, 2.0)
 0.0
 >>> foo(4, 5, 'loooong')
 26
 >>> foo(3, 4.5)
 Traceback (most recent call last):
 ...
 KeyError: (<class 'int'>, <class 'float'>)
 """
 @wraps(for_function)
 def decorated(*args):
 return decorated.registry[tuple(type(i) for i in args)](*args)
 decorated.registry = {}
 def adder(func):
 """
 Adds the supplied function to the multiple dispatch registry.
 """
 code = func.__code__
 args = code.co_varnames
 annotations = func.__annotations__
 types = tuple(annotations[a] for a in args if a in annotations)
 decorated.registry[types] = func
 return func
 decorated.add = adder
 return decorated

I'd like to read your comments about the code (efficiency, re-usability, etc.), documentation and everything else that you feel is worth noting.

Update: After fixing all the points listed in the choosen answer, I have created a new version.

Answer 1

1. You have a docstring and some doctests, which is great! The docstring needs some work, though, as described below.

2. The docstring needs to mention that the returned function has an add method, and explain what this method does.

3. The docstring needs to explain how the multiple dispatch works: that is, by exact match on the types of the arguments against the types of the annotations of the function instances. You ought to make it clear that arguments without annotations are ignored by the dispatcher.

4. The docstring says, "The returned function doesn't allow keyword arguments" but you might want to elaborate on exactly what this means.

5. The example code in the docstring is not very motivating. Why would I want to want to write a function foo like this? It would be very helpful to the reader if you could come up with an example that more closely resembles a real use case.

6. A group of multi-dispatch functions is a persistent data structure with two methods (add and __call__) and so would be clearer, I think, if it were implemented as a class.

7. The introspection code would be much clearer if you used inspect.signature. (This was introduced in 3.3, but since you're using function annotations, you probably don't mind this.)

8. The function that is decorated with @multi_dispatch is only used for its name, module, and docstring (the items that are copied by functools.update_wrapper). The function body is ignored. This seems like a waste: why not add this function to the registry too?

9. If you wrote return decorated instead of return func at the end of adder, then you wouldn't have to respell the names of the method instances.

   Update: What I mean by this is that you have to spell your method instances _foo rather than foo, to avoid overwriting the multi-method in the variable foo. With my suggestion, you wouldn't have to do this.

10. Dispatching on the exact types of the arguments limits the usefulness of multiple dispatch. In languages like Common Lisp that have multiple dispatch, dispatching can be performed on subclass matches (not just on exact matches). This makes for more natural fallback code, for example:

    @find.add
    def find(needle: str, haystack: str):
     # When both arguments are strings, Python has a built-in efficient
     # search operation.
     return haystack.find(needle)
    @find.add
    def find(needle: Sequence, haystack: Sequence):
     # Otherwise, for generic sequences, use Boyer–Moore–Horspool.
     h = len(haystack)
     n = len(needle)
     skip = {needle[i]: n - i - 1 for i in range(n - 1)}
     i = n - 1
     while i < h:
     for j in range(n):
     if haystack[i - j] != needle[-j - 1]:
     i += skip.get(haystack[i], n)
     break
     else:
     return i - n + 1
     return -1
    

    With your implementation one would have to provide separate method instances for (list, list), (tuple, tuple), (list, tuple) and other combinations of sequence types.

    Update: In comments you express concern about the speed of doing the subclass lookups. But you can deal with that by maintaining a cache of the actual argument types and the method instance that was found corresponding to those types. The revised dispatch logic would be something like this:

    arg_types = tuple(type(i) for i in args)
    try:
     method = cache[arg_types]
    except KeyError:
     method = ... # slow subclass dispatch logic here
     cache[arg_types] = method
    return method(*args)
    

    (When adding a new method instance, this might invalidate some of the cached lookups, so you'd need to clear the cache in the add method.)

11. You probably want to look at Guido van Rossum's multimethod implementation for other ideas (like "stacking" of multimethod decorators).

• Joschua: Thanks, good points! 9. I don't understand what you mean with respell the names of method instances. Couldn't I remove the return completely without changing anything, because every func is in the registry? 10. I agree, but wouldn't that make lookup much slower? 11. Actually, I've been inspired by his implementation. Maybe I'll do an (additional) alternative class that will both allow subclass matching and multiple of those.
• Joschua: Thanks again. If you also want to review my revised code: here is the link :)