¿Cómo paso argumentos adicionales a un decorador de Python?

100

Tengo un decorador como el de abajo.

def myDecorator(test_func):
    return callSomeWrapper(test_func)
def callSomeWrapper(test_func):
    return test_func
@myDecorator
def someFunc():
    print 'hello'

Quiero mejorar este decorador para aceptar otro argumento como el siguiente

def myDecorator(test_func,logIt):
    if logIt:
        print "Calling Function: " + test_func.__name__
    return callSomeWrapper(test_func)
@myDecorator(False)
def someFunc():
    print 'Hello'

Pero este código da el error,

TypeError: myDecorator () toma exactamente 2 argumentos (1 dado)

¿Por qué la función no se pasa automáticamente? ¿Cómo paso explícitamente la función a la función decoradora?

balki
fuente
3
balki: evite usar boolean como su argumento, no es un enfoque de gd y reduce la facilidad de lectura del código
Kit Ho
8
@KitHo: es una bandera booleana, por lo que usar un valor booleano es el enfoque correcto.
AKX
2
@KitHo - ¿Qué es "gd"? Esta bien"?
Rob Bednark

Respuestas:

173

Dado que está llamando al decorador como una función, necesita devolver otra función que es el decorador real:

def my_decorator(param):
    def actual_decorator(func):
        print("Decorating function {}, with parameter {}".format(func.__name__, param))
        return function_wrapper(func)  # assume we defined a wrapper somewhere
    return actual_decorator

La función externa recibirá cualquier argumento que pase explícitamente y debería devolver la función interna. A la función interna se le pasará la función para decorar y devolverá la función modificada.

Por lo general, desea que el decorador cambie el comportamiento de la función envolviéndola en una función contenedora. Aquí hay un ejemplo que opcionalmente agrega registro cuando se llama a la función:

def log_decorator(log_enabled):
    def actual_decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            if log_enabled:
                print("Calling Function: " + func.__name__)
            return func(*args, **kwargs)
        return wrapper
    return actual_decorator

La functools.wrapsllamada copia cosas como el nombre y la cadena de documentos en la función contenedora, para que sea más similar a la función original.

Uso de ejemplo:

>>> @log_decorator(True)
... def f(x):
...     return x+1
...
>>> f(4)
Calling Function: f
5
interjay
fuente
11
Y functools.wrapses recomendable usarlo : conserva el nombre original, la cadena de documentos, etc. de la función envuelta.
AKX
@AKX: Gracias, agregué esto al segundo ejemplo.
interjay
1
Entonces, básicamente, el decorador siempre toma solo un argumento, que es la función. Pero el decorador puede ser un valor de retorno de una función que podría aceptar argumentos. ¿Es esto correcto?
balki
2
@balki: Sí, eso es correcto. Lo que confunde las cosas es que muchas personas también llamarán myDecoratordecorador a la función externa ( aquí). Esto es conveniente para un usuario del decorador, pero puede resultar confuso cuando intentas escribir uno.
interjay
1
Pequeño detalle que me confundió: si log_decoratortoma un argumento predeterminado, no puede usarlo @log_decorator, debe serlo@log_decorator()
Stardidi
46

Solo para proporcionar un punto de vista diferente: la sintaxis

@expr
def func(...): #stuff

es equivalente a

def func(...): #stuff
func = expr(func)

En particular, exprpuede ser cualquier cosa que desee, siempre que se evalúe como un invocable. En particular en particular, exprpuede ser una fábrica de decoradores: le das unos parámetros y te da un decorador. Entonces, tal vez una mejor manera de comprender su situación es

dec = decorator_factory(*args)
@dec
def func(...):

que luego se puede acortar a

@decorator_factory(*args)
def func(...):

Por supuesto, dado que parece que decorator_factoryes un decorador, la gente tiende a nombrarlo para reflejar eso. Lo que puede resultar confuso cuando intentas seguir los niveles de indirección.

Katriel
fuente
Gracias, esto realmente me ayudó a entender la razón de ser de lo que está pasando.
Aditya Sriram
26

Solo quiero agregar algún truco útil que permita hacer que los argumentos del decorador sean opcionales. También permitirá reutilizar el decorador y reducir el anidamiento.

import functools

def myDecorator(test_func=None,logIt=None):
    if not test_func:
        return functools.partial(myDecorator, logIt=logIt)
    @functools.wraps(test_func)
    def f(*args, **kwargs):
        if logIt==1:
            print 'Logging level 1 for {}'.format(test_func.__name__)
        if logIt==2:
            print 'Logging level 2 for {}'.format(test_func.__name__)
        return test_func(*args, **kwargs)
    return f

#new decorator 
myDecorator_2 = myDecorator(logIt=2)

@myDecorator(logIt=2)
def pow2(i):
    return i**2

@myDecorator
def pow3(i):
    return i**3

@myDecorator_2
def pow4(i):
    return i**4

print pow2(2)
print pow3(2)
print pow4(2)
Alexey Smirnov
fuente
16

Solo otra forma de hacer decoradores. Encuentro que esta es la forma más fácil de entender.

class NiceDecorator:
    def __init__(self, param_foo='a', param_bar='b'):
        self.param_foo = param_foo
        self.param_bar = param_bar

    def __call__(self, func):
        def my_logic(*args, **kwargs):
            # whatever logic your decorator is supposed to implement goes in here
            print('pre action baz')
            print(self.param_bar)
            # including the call to the decorated function (if you want to do that)
            result = func(*args, **kwargs)
            print('post action beep')
            return result

        return my_logic

# usage example from here on
@NiceDecorator(param_bar='baaar')
def example():
    print('example yay')


example()
Robert Fey
fuente
¡gracias! He estado buscando algunas "soluciones" alucinantes durante unos 30 minutos y esta es la primera que realmente tiene sentido.
Canhazbits
0

Ahora, si desea llamar a una función function1con un decorador decorator_with_argy en este caso tanto la función como el decorador toman argumentos,

def function1(a, b):
    print (a, b)

decorator_with_arg(10)(function1)(1, 2)
SuperNova
fuente