Easy shimming in Python

Suppose you want to create a Python wrapper for some module, but hook into EVERY function call to do something with the args and results.

In my case, the module was a C Extension (using cffi), and I needed to convert most arguments to <cdata> things for cffi, and then turn most of the results back into Python objects.

Indeed, you could do a billion lines of this.

def my_func(a,b,c):
    return foo(lib.my_func(bar(a), bar(b), bar(c))

But nah, time to stir some metaprogramming into the soup.

I’ll assume you have functions like this in the wrapper module.

def do_something_with_arg(a):
    pass

def do_something_with_result(r):
    pass

You’ll need a decorator to actually apply the functions (though we won’t use it like a decorator). Pretty simple, just passes args and kwargs through one function, and the return value through the other. We use functools.wraps to preserve the original function’s name, docstring, etc.

def cast_decorator(f):
    """ considered calling this deCASTrator() but... """
    @functools.wraps(f)
    def wrapper(*args, **kwargs):
        for k, v in kwargs.iteritems():
            kwargs[k] = do_something_with_arg(v)
        args = map(do_something_with_arg, args)
        return do_something_with_result(f(*args, **kwargs))
    return wrapper

Now, we set the the module that the decorated members should be imported into. I wanted to dump all the resulting objects into the wrapper module (as if we did from SOMETHING.lib import *), so I just used __name__, but if that’s not the case for you, you’ll have to change this to another module.

current_module = __import__(__name__)

We iterate through all the objects in the source module, and apply the decorator if it’s a C function (change BuiltinFunctionType to FunctionType to process Python functions). Then, using setattr, we assign the hooked object to the target module with the original name.

from somewhere import lib
for name in dir(lib):
    obj = getattr(lib, name)
    if isinstance(obj, types.BuiltinFunctionType): # or types.FunctionType for non-native code
        if hasattr(obj, '__module__'):
            obj.__module__ = __name__ # for Sphinx to include it in autodoc
        setattr(current_module, name, cast_decorator(obj)) # apply decorator
    else:
        setattr(current_module, name, obj) # otherwise just straight up import it

That’s it, enjoy your meal.