2010
Aug
07

More Python voodoo: optional-argument decorators

I've just been doing something that probably should never be done in Python (again), and I figured it might be useful to record this little snippet for posterity. Many Python programmers are probably familiar with decorators, which modify functions by wrapping them with other functions. For example, if you want to print 'entering' and 'exiting' to trace when a function's execution begins and ends, you could do that with a decorator:

def trace(func):
    def wrapper(*args, **kwargs):
        print 'entering'
        func(*args, **kwargs)
        print 'exiting'
    return wrapper

@trace
def f(a,b):
     print 'in f: ' + str(a) + ',' + str(b)

Note that in practice there are better ways to do this, but it's just an example.

Anyway, it's also possible to customize the behavior of decorator itself by providing arguments, for instance if you wanted to print different strings instead of 'entering' and 'exiting'. But in that case, your "decorator" function is actually a wrapper that creates the real decorator and returns it:

def trace(enter_string, exit_string):
    def _trace(func):
        def wrapper(*args, **kwargs):
            print enter_string
            func(*args, **kwargs)
            print exit_string
        return wrapper
    return _trace

@trace('calling f', 'returning')
def f(a,b):
     print 'in f: ' + str(a) + ',' + str(b)

OK, so far so good. Slightly mind-bending, but that's Python for you.

Now, what if you want both ways to work? What if you want a decorator that can be used with or without arguments, as either

@trace
def f(a,b):
     print 'in f: ' + str(a) + ',' + str(b)

or

@trace('calling f', 'returning')
def f(a,b):
     print 'in f: ' + str(a) + ',' + str(b)

Well, in the former case the decorator will be called with one argument, the function to be decorated. You can identify this case in an if statement and treat it separately. Like so:

def trace(*args):
    def _trace(func):
        def wrapper(*args, **kwargs):
            print enter_string
            func(*args, **kwargs)
            print exit_string
        return wrapper
    if len(args) == 1 and callable(args[0]):
        # No arguments, this is the decorator
        # Set default values for the arguments
        enter_string = 'entering'
        exit_string = 'exiting'
        return _trace(args[0])
    else:
        # This is just returning the decorator
        enter_string, exit_string = args
        return _trace

In this example, _trace is the "real" decorator, the thing that gets called to wrap f. The top-level function trace either stores its arguments and returns the decorator itself (if it has arguments), or calls the decorator to do its job and returns the result (if there are no arguments provided, other than the function to be wrapped of course). Obviously, this won't work if you have a decorator that needs to take one argument (not the function it's wrapping) which is itself a callable function — in that case you're on your own.

In fact — if you're not confused yet — you can even make a decorator that turns your decorator functions into optional-argument decorators. A meta-decorator, I guess. Here's one way to do it:

def opt_arguments(func):
    def meta_wrapper(*args, **kwargs):
        if len(args) == 1 and callable(args[0]):
            # No arguments, this is the decorator
            # Set default values for the arguments
            return func(args[0])
        else:
            def meta_func(inner_func):
                return func(inner_func, *args, **kwargs)
            return meta_func
    return meta_wrapper

You can use this to decorate a decorator function which has one required argument and default values for all other arguments:

@opt_arguments
def trace(func, enter_string='entering', exit_string='exiting'):
    def wrapper(*args, **kwargs):
        print enter_string
        func(*args, **kwargs)
        print exit_string
    return wrapper

@trace
def f(a,b):
    print 'in f(%s,%s)' % (a,b)

@trace('call g', 'return g')
def g(a,b):
    print 'in g(%s,%s)' % (a,b)

By the way, I haven't tried this when the decorator is a class or callable object instead of a function, but I think it should work.

2009
Jul
27

More Python voodoo: combination class/instance methods

Here's a neat tidbit from my website code: Let's say you're writing a class, and you want it to have some methods that can be called as either class or instance methods. Of course, you can do it by writing a class method:

class F(object):
    @classmethod
    def foo(cls, a, b, c):
        print a, b, c

F.foo(1, 2, 3)
F().foo(1, 2, 3)

But what if you want a pattern like the following instead, where the method, when called on an instance, can use the instance's attributes?

F.foo(1, 2, 3)
F(1, 2, 3).foo()

You can do this with a descriptor. Descriptors allow you to customize the way attributes are accessed (and set); you can define precisely what you want F.foo or F().foo to do. This is kind of like the __getattribute__ method, but more general because you can use descriptors for class attributes as well.

class FooDescriptor(object):
    def __init__(self, method_name):
        super(FooDescriptor, self).__init__()
        self.method_name = method_name

    def __get__(self, instance, owner):
        if instance is None:
            # As a class attribute
            def foo(cls, a, b, c):
                print a, b, c
        else:
            # As an instance attribute
            def foo(self):
                print self.a, self.b, self.c
        return foo

The downside of this procedure is that you have to write two different definitions of the method, although that's generally only a minor inconvenience. If you're using this technique for a long method, it's not hard to split out most of the code into a common function that gets called from both the class method and instance method versions.