Como faço para passar argumentos extras para um decorador Python?

93

Eu tenho um decorador como abaixo.

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

Eu quero melhorar este decorador para aceitar outro argumento como abaixo

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

Mas este código dá o erro,

TypeError: myDecorator () recebe exatamente 2 argumentos (1 fornecido)

Por que a função não é passada automaticamente? Como faço para passar explicitamente a função para a função de decorador?

Balki
fonte
3
balki: evite usar booleano como seu argumento, não é uma abordagem gd e reduza a legibilidade do código
Kit Ho
8
@KitHo - é um sinalizador booleano, portanto, usar um valor booleano é a abordagem certa.
AKX de
2
@KitHo - o que é "gd"? Isso é bom"?
Rob Bednark

Respostas:

164

Como você está chamando o decorador como uma função, ele precisa retornar outra função que é o 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

A função externa receberá quaisquer argumentos que você passar explicitamente e deverá retornar a função interna. A função interna receberá a função a ser decorada e retornará a função modificada.

Normalmente, você deseja que o decorador altere o comportamento da função envolvendo-a em uma função de wrapper. Aqui está um exemplo que opcionalmente adiciona registro quando a função é chamada:

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

A functools.wrapschamada copia coisas como o nome e docstring para a função de wrapper, para torná-la mais semelhante à função original.

Exemplo de uso:

>>> @log_decorator(True)
... def f(x):
...     return x+1
...
>>> f(4)
Calling Function: f
5
interjay
fonte
11
E o uso functools.wrapsé aconselhável - ele retém o nome original, docstring, etc. da função agrupada.
AKX de
@AKX: Obrigado, adicionei isso ao segundo exemplo.
intervalo de
1
Então, basicamente, o decorator sempre leva apenas um argumento que é a função. Mas o decorador pode ser um valor de retorno de uma função que pode receber argumentos. Isso está correto?
balki
2
@balki: Sim, correto. O que confunde as coisas é que muitas pessoas também chamarão a função externa ( myDecoratoraqui) de decorador. Isso é conveniente para um usuário do decorador, mas pode ser confuso quando você está tentando escrever um.
intervalo de
1
Pequeno detalhe que me confundiu: se você log_decoratorpega um argumento padrão, você não pode usar @log_decorator, deve ser@log_decorator()
Stardidi
44

Apenas para fornecer um ponto de vista diferente: a sintaxe

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

é equivalente a

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

Em particular, exprpode ser o que você quiser, contanto que seja avaliado como um chamável. Em particular , exprpode ser uma fábrica de decoradores: você fornece alguns parâmetros e ela fornece um decorador. Então, talvez a melhor maneira de entender sua situação seja como

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

que pode então ser encurtado para

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

Claro, uma vez que parece como decorator_factoryum decorador, as pessoas tendem a nomeá-lo para refletir isso. O que pode ser confuso quando você tenta seguir os níveis de indireção.

Katriel
fonte
Obrigado, isso realmente me ajudou a entender a razão por trás do que está acontecendo.
Aditya Sriram
24

Só quero adicionar alguns truques úteis que permitirão tornar opcionais os argumentos do decorador. Também permitirá reutilizar o decorador e diminuir o aninhamento

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
fonte
16

Apenas outra maneira de fazer decoradores. Acho que essa maneira é a mais fácil de envolver minha cabeça.

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
fonte
obrigado! Estou procurando por algumas "soluções" alucinantes por cerca de 30 minutos e esta é a primeira que realmente faz sentido.
canhazbits
0

Agora, se você quiser chamar uma função function1com um decorador decorator_with_arge, neste caso, a função e o decorador usam argumentos,

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

decorator_with_arg(10)(function1)(1, 2)
Super Nova
fonte