Easy shimming in Python
AKA hooking the args & return value of every function in a Python module
23 Jun 2017Suppose 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.
