web-dev-qa-db-fra.com

Comment passer des arguments supplémentaires à un décorateur Python?

J'ai un décorateur comme ci-dessous.

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

Je veux améliorer ce décorateur pour accepter un autre argument comme ci-dessous

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

Mais ce code donne l'erreur,

TypeError: myDecorator () prend exactement 2 arguments (1 donné)

Pourquoi la fonction n'est-elle pas transmise automatiquement? Comment passer explicitement la fonction à la fonction décorateur?

78
balki

Puisque vous appelez le décorateur comme une fonction, il doit renvoyer une autre fonction qui est le décorateur réel:

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 fonction externe recevra tous les arguments que vous passez explicitement et devrait renvoyer la fonction interne. La fonction interne recevra la fonction à décorer et renverra la fonction modifiée.

Habituellement, vous souhaitez que le décorateur modifie le comportement de la fonction en l'enveloppant dans une fonction wrapper. Voici un exemple qui ajoute éventuellement la journalisation lorsque la fonction est appelée:

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

Le functools.wraps call copie des choses comme le nom et la docstring dans la fonction wrapper, pour la rendre plus similaire à la fonction d'origine.

Exemple d'utilisation:

>>> @log_decorator(True)
... def f(x):
...     return x+1
...
>>> f(4)
Calling Function: f
5
139
interjay

Juste pour fournir un point de vue différent: la syntaxe

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

est équivalent à

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

En particulier, expr peut être tout ce que vous aimez, tant qu'il est évalué comme appelable. Dans particulier particulier, expr peut être une fabrique de décorateur: vous lui donnez quelques paramètres et cela vous donne un décorateur. Alors peut-être qu'une meilleure façon de comprendre votre situation est

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

qui peut ensuite être raccourci en

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

Bien sûr, car il ressemble comme decorator_factory est décorateur, les gens ont tendance à le nommer pour refléter cela. Ce qui peut être déroutant lorsque vous essayez de suivre les niveaux d'indirection.

41
Katriel

Je veux juste ajouter une astuce utile qui permettra de rendre les arguments du décorateur facultatifs. Il permettra également de réutiliser le décorateur et de diminuer l'imbrication

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)
17
Alexey Smirnov

Juste une autre façon de faire des décorateurs. Je trouve que c'est la façon la plus facile d'envelopper ma tête.

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()
7
Robert Fey