decorations.py

from functools import wraps, partial
from inspect import getargspec

def _param_update(old_params, new_params):
    args, kwds = old_params

    if type(new_params) == tuple:
        if len(new_params) == 2 and type(new_params[0]) == tuple and type(new_params[1]) == dict:
            args = new_params[0]
            kwds.update(new_params[1])
        else:
            args = new_params
    elif type(new_params) == dict:
        kwds.update(new_params)

    return (args, kwds)


def _metadecorator(decoration, params=False):
    "Takes a decoration to pre-apply to a decoratee"
    def cap_params(*dec_params, **dec_kwds):
        "Captures decoration parameters"
        # print decoration, params, dec_params, dec_kwds
        @wraps(decoration)
        def decorator(decoratee):
            "The decorator - takes a decoratee to decorate"
            @wraps(decoratee)
            def decorated(*args, **kwds):
                "The decorated decoratee"
                # the decoration will be passed the decoratee and (optionally) two sets of arguments
                decoration_args = {
                    'dec_params': dec_params,
                    'dec_kwds': dec_kwds,
                    'args': args,
                    'kwds': kwds
                }
                filtered_args = dict( (arg, decoration_args.get(arg)) for arg in getargspec(decoration).args[1:] )
                return decoration(decoratee, **filtered_args)
            return decorated
            
        if params=="maybe" and not dec_kwds and len(dec_params)==1 and callable(dec_params[0]):
            decoratee = dec_params[0]
            dec_params = ()
            return decorator(decoratee)
        else:
            return decorator
            
    return cap_params

def decoration(decoration=None, has_params="maybe"):
    "Meta-decorates"
    if decoration:
        dec_args = getargspec(decoration).args
        if set(dec_args).intersection(('dec_params', 'dec_kwds')):
            return _metadecorator(decoration, params=has_params)
        else:
            return _metadecorator(decoration, params=has_params)()
    else:
        return partial(decoration, params=has_params)

@decoration
def predecoration(f, args, kwargs):
    predecoratee = args[0]
    @wraps(predecoratee)
    def decoratee(*args, **kwargs):
        args, kwargs = _param_update((args, kwargs), f(*(args or []), **(kwargs or {})))
        return predecoratee(*(args or []), **(kwargs or {}))
    return decoratee

@decoration
def postdecoration(f, args, kwargs):
    postdecoratee = args[0]
    @wraps(postdecoratee)
    def decoratee(*args, **kwargs):
        result = postdecoratee(*(args or []), **(kwargs or {}))
        return f(result)
    return decoratee

# -- examples --

@predecoration
def debug_input(*input):
    print "ARGS WERE:", input
    
@postdecoration
def debug_output(output):
    print "RETVAL WAS:", output
    return output

@debug_input
@debug_output
def foo(x,y,z):
    return x+y

foo(1,2,3)

# prints:
#   ARGS WERE: (1, 2, 3)
#   RETVAL WAS: 3