Easy Python Decorators with the decorator decorator

Decorators in python are very powerful, but are often a pain to get right, especially when you want to pass arguments to the decorator. Typically, it involves defining a function in a function in a function, etc, up to 4 layers deep. Can we make it easier? What if we even made a decorator to make decorators?

Perhaps we could use it catch and log errors:

@decorator
def log_error(func, *args, **kargs):
   try:
       return func(*args, **kargs)
   except Exception, error:
       print "error occurred: %r" % error

@log_error
def blowup():
   raise Exception("blew up")

blowup() # this gets caught and prints the error

Or maybe we'd like to synchronize a function to make it thread-safe:

@decorator
def synchronized(lock, func, *args, **kargs):
    lock.acquire()
    try:
        return func(self, *args, **kargs)
    finally:
        lock.release()

missle_lock = thread.RLock()

@synchronized(missle_lock)
def launch_missiles():
    pass

Or we could even use arguments to do do some object __init__ trickery (something I've had to do when working with wxpython):

@decorator
def inherits_format(method, *args, **kargs):
    self, parent = args[:2]
    self.format = parent.format
    return method(*args, **kargs)

class Child:
    @inherits_format
    def __init__(self, parent):
        pass

class Parent:
    def __init__(self, format):
        self.format = format

format = object()
child  = Child(Parent(format))
assert child.format is format

If you've never worked python decorators, these are few examples of how powerful they are. But you'll probably never want to write your own because it can be a pain. If that's the case, then the decorator decorator is for you! Here's the code for it. It's a little tricky, but all you have to do is import it, slap @decorator before your decorator, make sure you call func(*args, **kargs), and you're all set. This is as easy as decorators can get.

Update: Dave left a comment and notified me that there is another, much more complete, implementation of the same idea at http://www.phyast.pitt.edu/~micheles/python/documentation.html. So, if you want a very complete module, go there. If you want something small and simple to cut and paste into your own code, use the following.

def decorator(decorator_func):
    decorator_expected_arg_count = decorator_func.func_code.co_argcount - 1
    if decorator_expected_arg_count:
        def decorator_maker(*decorator_args):
            assert len(decorator_args) == decorator_expected_arg_count, "%s expects %d args" % (decorator_func.__name__, decorator_expected_arg_count)

            def _decorator(func):
                assert callable(func), "Decorator not given a function.  Did you forget to give %r arguments?" % (decorator_func.__name__)

                def decorated_func(*args, **kargs):
                    full_args = decorator_args + (func,) + args
                    return decorator_func(*full_args, **kargs)

                decorated_func.__name__ = func.__name__
                return decorated_func
            return _decorator
        return decorator_maker
    else:
        def _decorator(func):
            def decorated_func(*args, **kargs):
                return decorator_func(func, *args, **kargs)

            decorated_func.__name__ = func.__name__
            return decorated_func
        return _decorator