Recipe20.6.Adding Functionality to a Class by Wrapping a Method

Recipe 20.6. Adding Functionality to a Class by Wrapping a Method

Credit: Ken Seehof, Holger Krekel

Problem

You need to add functionality to an existing class, without changing the source code for that class, and inheritance is not applicable (since it would make a new class, rather than changing the existing one). Specifically, you need to enrich a method of the class, adding some extra functionality "around" that of the existing method.

Solution

Adding completely new methods (and other attributes) to an existing class object is quite simple, since the built-in function setattr does essentially all the work. We need to "decorate" an existing method to add to its functionality. To achieve this, we can build the new replacement method as a closure. The best architecture is to define general-purpose wrapper and unwrapper functions, such as:

import inspect def wrapfunc(obj, name, processor, avoid_doublewrap=True):     """ patch obj.<name> so that calling it actually calls, instead,             processor(original_callable, *args, **kwargs)     """     # get the callable at obj.<name>     call = getattr(obj, name)     # optionally avoid multiple identical wrappings     if avoid_doublewrap and getattr(call, 'processor', None) is processor:         return     # get underlying function (if any), and anyway def the wrapper closure     original_callable = getattr(call, 'im_func', call)     def wrappedfunc(*args, **kwargs):         return processor(original_callable, *args, **kwargs)     # set attributes, for future unwrapping and to avoid double-wrapping     wrappedfunc.original = call     wrappedfunc.processor = processor     # 2.4 only: wrappedfunc._ _name_ _ = getattr(call, '_ _name_ _', name)     # rewrap staticmethod and classmethod specifically (iff obj is a class)     if inspect.isclass(obj):         if hasattr(call, 'im_self'):             if call.im_self:                 wrappedfunc = classmethod(wrappedfunc)         else:             wrappedfunc = staticmethod(wrappedfunc)     # finally, install the wrapper closure as requested     setattr(obj, name, wrappedfunc) def unwrapfunc(obj, name):     ''' undo the effects of wrapfunc(obj, name, processor) '''     setattr(obj, name, getattr(obj, name).original)

This approach to wrapping is carefully coded to work just as well on ordinary functions (when obj is a module) as on methods of all kinds (e.g., bound methods, when obj is an instance; unbound, class, and static methods, when obj is a class). This method doesn't work when obj is a built-in type, though, because built-ins are immutable.

For example, suppose we want to have "tracing" prints of all that happens whenever a particular method is called. Using the general-purpose wrapfunc function just shown, we could code:

def tracing_processor(original_callable, *args, **kwargs):     r_name = getattr(original_callable, '_ _name_ _', '<unknown>')     r_args = map(repr, args)     r_args.extend(['%s=%r' % x for x in kwargs.iteritems( )])     print "begin call to %s(%s)" % (r_name, ", ".join(r_args))     try:         result = call(*args, **kwargs)     except:         print "EXCEPTION in call to %s" %(r_name,)         raise     else:         print "call to %s result: %r" %(r_name, result)         return result def add_tracing_prints_to_method(class_object, method_name):     wrapfunc(class_object, method_name, tracing_processor)

Discussion

This recipe's task occurs fairly often when you're trying to modify the behavior of a standard or third-party Python module, since editing the source of the module itself is undesirable. In particular, this recipe can be handy for debugging, since the example function add_tracing_prints_to_method presented in the "Solution" lets you see on standard output all details of calls to a method you want to watch, without modifying the library module, and without requiring interactive access to the Python session in which the calls occur.

You can also use this recipe's approach on a larger scale. For example, say that a library that you imported has a long series of methods that return numeric error codes. You could wrap each of them inside an enhanced wrapper method, which raises an exception when the error code from the original method indicates an error condition. Again, a key issue is not having to modify the library's own code. However, methodical application of wrappers when building a subclass is also a way to avoid repetitious code (i.e., boilerplate). For example, Recipe 5.12 and Recipe 1.24 might be recoded to take advantage of the general wrapfunc presented in this recipe.

Particularly when "wrapping on a large scale", it is important to be able to "unwrap" methods back to their normal state, which is why this recipe's Solution also includes an unwrapfunc function. It may also be handy to avoid accidentally wrapping the same method in the same way twice, which is why wrapfunc supports the optional parameter avoid_doublewrap, defaulting to true, to avoid such double wrapping. (Unfortunately, classmethod and staticmethod do not support per-instance attributes, so the avoidance of double wrapping, as well as the ability to "unwrap", cannot be guaranteed in all cases.)

You can wrap the same method multiple times with different processors. However, unwrapping must proceed last-in, first-out; as coded, this recipe does not support the ability to remove a wrapper from "somewhere in the middle" of a chain of several wrappers. A related limitation of this recipe as coded is that double wrapping is not detected when another unrelated wrapping occurred in the meantime. (We don't even try to detect what we might call "deep double wrapping.")

If you need "generalized unwrapping", you can extend unwrap_func to return the processor it has removed; then you can obtain generalized unwrapping by unwrapping all the way, recording a list of the processors that you removed, and then pruning that list of processors and rewrapping. Similarly, generalized detection of "deep" double wrapping could be implemented based on this same idea.

Another generalization, to fully support staticmethod and classmethod, is to use a global dict, rather than per-instance attributes, for the original and processor values; functions, bound and unbound methods, as well as class methods and static methods, can all be used as keys into such a dictionary. Doing so obviates the issue with the inability to set per-instance attributes on class methods and static methods. However, each of these generalizations can be somewhat complicated, so we are not pursuing them further here.

Once you have coded some processors with the signature and semantics required by this recipe's wrapfunc, you can also use such processors more directly (in cases where modifying the source is OK) with a Python 2.4 decorator, as follows:

def processedby(processor):     """ decorator to wrap the processor around a function. """     def processedfunc(func):         def wrappedfunc(*args, **kwargs):             return processor(func, *args, **kwargs)         return wrappedfunc     return processedfunc

For example, to wrap this recipe's tracing_processor around a certain method at the time the class statement executes, in Python 2.4, you can code:

class SomeClass(object):     @processedby(tracing_processor)     def amethod(self, s):         return 'Hello, ' + s

Recipe20.6.Adding Functionality to a Class by Wrapping a Method