Question

What is the proper way in Python to allow the user to extend the types on which a function can operate without altering the original code of the function?

Suppose I have a module with a my_module.foo() function that was originally written to work on float types. Now I would like the same function to be able to work also with, let's say, mpmath arbitrary precision floats - but without altering the code in the original module.

In C++ I would add an extra overload (or, more likely, some template specialisation trickery with an helper struct). How should I structure the original my_module.foo() code so that the user can add her own custom hooks inside?

I can think of a few way to achieve this, but as a novice Python programmer I am sure most of them are going to be horrid :)

EDIT: thanks for all the answers so far, much appreciated.

I should probably clarify that one key requirement is being able to cope with types which I have not defined myself. E.g., if I am trying to code a generic cos function in my module, I want to call math.cos on builtin types, mpmath.cos on mpf types, sympy.cos on sympy symbolic types, etc. And of course I would like the dispatching logic not to be in my module's cos implementation.

Was it helpful?

Solution

There are two ways to do this:

  • Delegation. Already in Python. Most Pythonic. Doesn't work for builtin types. Not exactly what you're looking for.
  • Single Dispatch. Still a PEP. Works for builtin types. Precisely what you're looking for.

Delegation

You usually delegate responsibility to the object you're acting on, and you don't implement the logic in your function.

Here's an example: len. The implementation of len is very staightforward:

def len(obj):
    return obj.__len__()

Different types (str, list, tuple...) have different implementations, but they all work with the same function.

Now, if I want to define my own type, that works with len, I can do:

class MyTypeOfLength3(object):
    def __len__(self):
        return 3

o = MyTypeOfLength3()
print len(o) 
# 3

In your case, you would be implementing something that looks like len.

(Note: this isn't the actual code for len, but it's more or less equivalent.)

Single Dispatch

Of course, in some cases, that might not be practical. If that's your case, then the "Single Dispatch" PEP 443 is probably what you're looking for.

It suggests a new decorator that would accomplish what you're looking for:

>>> from functools import singledispatch
>>> @singledispatch
... def fun(arg, verbose=False):
...     if verbose:
...         print("Let me just say,", end=" ")
...     print(arg)
...
>>> @fun.register(int)
... def _(arg, verbose=False):
...     if verbose:
...         print("Strength in numbers, eh?", end=" ")
...     print(arg)
...
>>> @fun.register(list)
... def _(arg, verbose=False):
...     if verbose:
...         print("Enumerate this:")
...     for i, elem in enumerate(arg):
...         print(i, elem)

Once you've defined your function as such, you can call fun(something), and Python will find out the correct implementation (int or list here), of fallback to the default implementation def fun(...): ....

You therefore solely have to decorate your original function, and you're done, your users can add their own types.

Note: as pointed out in the comments, singledispatch is already implemented in Python, it's pkgutil.simplegeneric

OTHER TIPS

It's possible to do what you want without waiting for PEP 443 - Single-dispatch generic functions to be implemented by instead using abstract base classes which were added in Python 2.6. These allow you to create "virtual" meta classes and add arbitrary subclasses to them on the fly without altering existing code or monkey-patching it. Your module can then use types registered with this metaclass to figure out what to do. You (or the authors) of other types can register them as needed.

Here's example code illustrating the concept:

import abc

class Trigonometric(object):
    __metaclass__ = abc.ABCMeta
    _registry = {}

    @classmethod
    def register(cls, subclass, cos_func, sin_func):
        cls.__metaclass__.register(cls, subclass)
        if subclass not in cls._registry:  # may or may not want this check...
            cls._registry[subclass] = {'cos': cos_func, 'sin': sin_func}

    @classmethod
    def call_func(cls, func_name, n):
        try:
            return cls._registry[n.__class__][func_name](n)
        except KeyError:
            raise RuntimeError(
                "Either type {} isn't registered or function {}() "
                "isn't known.".format(n.__class__.__name__, func_name))

# module-level functions
def cos(n):
    return Trigonometric.call_func('cos', n)

def sin(n):
    return Trigonometric.call_func('sin', n)

if __name__ == '__main__':
    # avoid hardcoding this module's filename into the source
    import sys
    my_module = sys.modules[__name__]  # replaces import my_module

    # register the built-in float type
    import math
    print 'calling Trigonometric.register(float)'
    Trigonometric.register(float, math.cos, math.sin)

    # register mpmath's arbitrary-precision mpf float type
    from mpmath import mp
    print 'calling Trigonometric.register(mp.mpf)'
    Trigonometric.register(mp.mpf, mp.cos, mp.sin)

    f = 1.0
    print 'isinstance(f, Trigonometric):', isinstance(f, Trigonometric)
    print 'my_module.cos(f):', my_module.cos(f), my_module.sin(f)

    v = mp.mpf(1)
    print 'isinstance(v, Trigonometric):', isinstance(v, Trigonometric)
    print 'my_module.cos(v):', my_module.cos(v), my_module.sin(v)

It depends on a code and expected results. Usually you shouldn't specify data types implicitly. Use Duck Typing.

But if a function expects only floats, you can wrap it. There's a simplest example:

def bar(data):
    """Execute my_module.foo function with data converted to float"""
    return my_module.foo(float(data))
Licensed under: CC-BY-SA with attribution
Not affiliated with StackOverflow
scroll top