15.10 Adding Functionality to a Class

Credit: Ken Seehof

15.10.1 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, not change the old one).

15.10.2 Solution

Again, this is a case for introspection and dynamic change. The enhance_method function alters a klass class object to substitute a named method with an enhanced version, decorated by the replacement function argument. The method_logger method exemplifies a typical case of replacement by decorating any method with print statements tracing its calls and returns:

# requires Python 2.1, or 2.2 with classic classes only

from _ _future_ _ import nested_scopes
import new

def enhance_method(klass, method_name, replacement):
    'replace a method with an enhanced version'
    method = getattr(klass, method_name)
    def enhanced(*args, **kwds): return replacement(method, *args, **kwds)
    setattr(klass, method_name, new.instancemethod(enhanced, None, klass))

def method_logger(old_method, self, *args, **kwds):
    'example of enhancement: log all calls to a method'
    print '*** calling: %s%s, kwds=%s' % (old_method._ _name_ _, args, kwds)
    return_value = old_method(self, *args, **kwds) # call the original method
    print '*** %s returns: %r' % (old_method._ _name_ _, return_value)
    return return_value

def demo(  ):
    class Deli:
        def order_cheese(self, cheese_type):
            print 'Sorry, we are completely out of %s' % cheese_type

    d = Deli(  )
    d.order_cheese('Gouda')

    enhance_method(Deli, 'order_cheese', method_logger)
    d.order_cheese('Cheddar')

15.10.3 Discussion

This recipe is useful when you need to modify the behavior of a standard or third-party Python module, but changing the module itself is undesirable. In particular, this recipe can be handy for debugging, since you can use it to log all calls to a library method that you want to watch without changing the library code or needing interactive access to the session. The method_logger function in the recipe shows this specific logging usage, and the demo function shows typical usage.

Here's another, perhaps more impressive, use for this kind of approach. Sometimes you need to globally change the behavior of an entire third-party Python library. For example, say a Python library that you downloaded has 50 different methods that all return error codes, but you want these methods to raise exceptions instead (again, you don't want to change their code). After importing the offending module, you repeatedly call this recipe's enhance_method function to hook a replacement version that checks the return value and issues an exception if an error occurred around each method, wrapping each of the 50 methods in question with the same enhancement metafunction.

The heart of the recipe is the enhance_method function, which takes the class object, method name string, and replacement decorator function as arguments. It extracts the method with the getattr built-in function and replaces the method with the reciprocal setattr built-in function. The replacement is a new instance method (actually, an unbound method, as specified by the second None argument to new.instancemethod) that wraps an enhanced function, which is built with a local def. This relies on lexically nested scopes, since the local (nested) enhanced function must be able to see the method and replacement names that are local variables of the enclosing (outer) enhance_method function. The reliance on nested scopes is the reason this recipe specifies Python 2.1 or 2.2 (to work in 2.1, it needs the from _ _future_ _ import nested_scopes statement at the start of the module).

15.10.4 See Also

Recipe 15.7; Recipe 5.14 and Recipe 15.11 for other approaches to modifying the methods of an instance; documentation on the new standard library module in the Library Reference.