How can I get a Python decorator to run after the decorated function has completed?

I want to use a decorator to handle auditing of various functions (mainly Django view functions, but not exclusively). In order to do this I would like to be able to audit the function post-execution - i.e. the function runs as normal, and if it returns without an exception, then the decorator logs the fact.

Something like:

@audit_action(action='did something')
def do_something(*args, **kwargs):
    if args[0] == 'foo':
        return 'bar'
    else:
        return 'baz'

Where audit_action would only run after the function has completed.

Solution

Decorators usually return a wrapper function; just put your logic in the wrapper function after invoking the wrapped function.

def audit_action(action):
    def decorator_func(func):
        def wrapper_func(*args, **kwargs):
            # Invoke the wrapped function first
            retval = func(*args, **kwargs)
            # Now do something here with retval and/or action
            print('In wrapper_func, handling action {!r} after wrapped function returned {!r}'.format(action, retval))
            return retval
        return wrapper_func
    return decorator_func

So audit_action(action='did something') is a decorator factory that returns a scoped decorator_func, which is used to decorate your do_something (do_something = decorator_func(do_something)).

After decorating, your do_something reference has been replaced by wrapper_func. Calling wrapper_func() causes the original do_something() to be called, and then your code in the wrapper func can do things.

The above code, combined with your example function, gives the following output:

>>> do_something('foo')
In wrapper_func, handling action 'did something' after wrapped function returned 'bar'
'bar'