Decoradores con parámetros?

Una forma de pensar acerca de los decoradores con argumentos es{[15]]}

@decorator
def foo(*args, **kwargs):
    pass

Se traduce como

foo = decorator(foo)

Así que si el decorador tenía argumentos,{[15]]}

@decorator_with_args(arg)
def foo(*args, **kwargs):
    pass

Se traduce como

foo = decorator_with_args(arg)(foo)

decorator_with_args es una función que acepta un argumento personalizado y que devuelve el decorador real (que se aplicará a la función decorada).

Uso un truco simple con parciales para hacer que mis decoradores sean fáciles

from functools import partial

def _pseudo_decor(fun, argument):
    def ret_fun(*args, **kwargs):
        #do stuff here, for eg.
        print ("decorator arg is %s" % str(argument))
        return fun(*args, **kwargs)
    return ret_fun

real_decorator = partial(_pseudo_decor, argument=arg)

@real_decorator
def foo(*args, **kwargs):
    pass

Actualización:

Anterior, foo se convierte en real_decorator(foo)

Un efecto de decorar una función es que el nombre foo se anula en la declaración del decorador. foo es "anulado" por lo que sea devuelto por real_decorator. En este caso, un nuevo objeto de función.

Todos los metadatos de foo están sobreescritos, notablemente docstring y function name.

>>> print(foo)
<function _pseudo_decor.<locals>.ret_fun at 0x10666a2f0>

Functools.wraps nos da un método conveniente para "levantar" el docstring y el nombre a la función devuelta.

from functools import partial, wraps

def _pseudo_decor(fun, argument):
    # magic sauce to lift the name and doc of the function
    @wraps(fun)
    def ret_fun(*args, **kwargs):
        #do stuff here, for eg.
        print ("decorator arg is %s" % str(argument))
        return fun(*args, **kwargs)
    return ret_fun

real_decorator = partial(_pseudo_decor, argument=arg)

@real_decorator
def bar(*args, **kwargs):
    pass

>>> print(bar)
<function __main__.bar(*args, **kwargs)>

Me gustaría mostrar una idea que en mi humilde opinión es bastante elegante. La solución propuesta por t.dubrownik muestra un patrón que es siempre el mismo: necesita la envoltura de tres capas independientemente de lo que haga el decorador.

Así que pensé que este es un trabajo para un meta-decorador, es decir, un decorador para decoradores. Como un decorador es una función, en realidad funciona como un decorador regular con argumentos:

def parametrized(dec):
    def layer(*args, **kwargs):
        def repl(f):
            return dec(f, *args, **kwargs)
        return repl
    return layer

Esto se puede aplicar a un decorador normal para agregar parámetros. Así que para ejemplo, digamos que tenemos el decorador que duplica el resultado de una función:

def double(f):
    def aux(*xs, **kws):
        return 2 * f(*xs, **kws)
    return aux

@double
def function(a):
    return 10 + a

print function(3)    # Prints 26, namely 2 * (10 + 3)

Con @parametrized podemos construir un decorador genérico @multiply que tenga un parámetro

@parametrized
def multiply(f, n):
    def aux(*xs, **kws):
        return n * f(*xs, **kws)
    return aux

@multiply(2)
def function(a):
    return 10 + a

print function(3)    # Prints 26

@multiply(3)
def function_again(a):
    return 10 + a

print function(3)          # Keeps printing 26
print function_again(3)    # Prints 39, namely 3 * (10 + 3)

Convencionalmente el primer parámetro de un decorador parametrizado es la función, mientras que los argumentos restantes corresponderán al parámetro del decorador parametrizado.

Un ejemplo de uso interesante podría ser un decorador asertivo seguro de tipo:

import itertools as it

@parametrized
def types(f, *types):
    def rep(*args):
        for a, t, n in zip(args, types, it.count()):
            if type(a) is not t:
                raise TypeError('Value %d has not type %s. %s instead' %
                    (n, t, type(a))
                )
        return f(*args)
    return rep

@types(str, int)  # arg1 is str, arg2 is int
def string_multiply(text, times):
    return text * times

print(string_multiply('hello', 3))    # Prints hellohellohello
print(string_multiply(3, 3))          # Fails miserably with TypeError

Una nota final: aquí no estoy usando functools.wraps para las funciones wrapper, pero yo recomendaría usarlo todas las veces.

Aquí hay una versión ligeramente modificada de la respuesta de t.dubrownik. ¿Por qué?

1. Como plantilla general, debe devolver el valor devuelto de la función original.
2. Esto cambia el nombre de la función, lo que podría afectar a otros decoradores / código.

Así que use @functools.wraps():

from functools import wraps

def decorator(argument):
    def real_decorator(function):
        @wraps(function)
        def wrapper(*args, **kwargs):
            funny_stuff()
            something_with_argument(argument)
            retval = function(*args, **kwargs)
            more_funny_stuff()
            return retval
        return wrapper
    return real_decorator

Supongo que su problema es pasar argumentos a su decorador. Esto es un poco complicado y no sencillo.

Aquí hay un ejemplo de cómo hacer esto:

class MyDec(object):
    def __init__(self,flag):
        self.flag = flag
    def __call__(self, original_func):
        decorator_self = self
        def wrappee( *args, **kwargs):
            print 'in decorator before wrapee with flag ',decorator_self.flag
            original_func(*args,**kwargs)
            print 'in decorator after wrapee with flag ',decorator_self.flag
        return wrappee

@MyDec('foo de fa fa')
def bar(a,b,c):
    print 'in bar',a,b,c

bar('x','y','z')

Impresiones:

in decorator before wrapee with flag  foo de fa fa
in bar x y z
in decorator after wrapee with flag  foo de fa fa

Vea el artículo de Bruce Eckel para más detalles.

En mi caso, decidí resolver esto a través de una lambda de una línea para crear una nueva función de decorador:

def finished_message(function, message="Finished!"):

    def wrapper(*args, **kwargs):
        output = function(*args,**kwargs)
        print(message)
        return output

    return wrapper

@finished_message
def func():
    pass

my_finished_message = lambda f: finished_message(f, "All Done!")

@my_finished_message
def my_func():
    pass

if __name__ == '__main__':
    func()
    my_func()

Cuando se ejecuta, esto imprime:

Finished!
All Done!

Quizás no tan extensible como otras soluciones, pero funcionó para mí.