Как передать дополнительные аргументы декоратору Python?

100

У меня есть декоратор, как показано ниже.

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

Я хочу улучшить этот декоратор, чтобы он принял еще один аргумент, как показано ниже

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

Но этот код дает ошибку,

TypeError: myDecorator () принимает ровно 2 аргумента (1 задан)

Почему функция не передается автоматически? Как явно передать функцию функции декоратора?

балки
источник
3
balki: пожалуйста, избегайте использования логического значения в качестве аргумента, это не подход gd и снижает удобочитаемость кода
Кит Хо
8
@KitHo - это логический флаг, поэтому использование логического значения - правильный подход.
AKX
2
@KitHo - что такое "б-г"? Это хорошо"?
Роб Беднарк

Ответы:

173

Поскольку вы вызываете декоратор как функцию, он должен вернуть другую функцию, которая является фактическим декоратором:

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

Внешняя функция получит любые аргументы, которые вы передадите явно, и должна вернуть внутреннюю функцию. Внутренняя функция будет передана функции для украшения и вернет измененную функцию.

Обычно вы хотите, чтобы декоратор изменил поведение функции, заключив ее в функцию-оболочку. Вот пример, который необязательно добавляет ведение журнала при вызове функции:

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

Эти functools.wrapsкопии вызова вещь , как имя и строка документация для функции обертки, чтобы сделать его более похожим на исходную функцию.

Пример использования:

>>> @log_decorator(True)
... def f(x):
...     return x+1
...
>>> f(4)
Calling Function: f
5
интерджей
источник
11
И functools.wrapsрекомендуется использовать - он сохраняет исходное имя, строку документации и т. Д. Обернутой функции.
AKX
@AKX: Спасибо, я добавил это ко второму примеру.
Interjay
1
Таким образом, декоратор всегда принимает только один аргумент - функцию. Но декоратор может быть возвращаемым значением функции, которая может принимать аргументы. Это правильно?
balki
2
@balki: Да, это правильно. Смущает то, что многие люди также называют внешнюю функцию ( myDecoratorздесь) декоратором. Это удобно для пользователя декоратора, но может сбивать с толку, когда вы пытаетесь его написать.
Interjay
1
Маленькая деталь, которая меня смутила: если вы log_decorator@log_decorator@log_decorator()
берете
46

Просто чтобы представить другую точку зрения: синтаксис

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

эквивалентно

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

В частности, exprможет быть что угодно, если оно оценивается как вызываемый. В частности , это exprможет быть фабрика декораторов: вы задаете ей некоторые параметры, и она дает вам декоратор. Так что, возможно, лучший способ понять вашу ситуацию - это

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

который затем можно сократить до

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

Конечно, поскольку он выглядит как decorator_factoryдекоратор, люди склонны называть его, чтобы отразить это. Что может сбивать с толку, когда вы пытаетесь следовать уровням косвенного обращения.

Катриэль
источник
Спасибо, это действительно помогло мне понять причину того, что происходит.
Адитья Шрирам
26

Просто хочу добавить полезный трюк, который позволит сделать аргументы декоратора необязательными. Это также позволит повторно использовать декоратор и уменьшить вложенность

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)
Алексей Смирнов
источник
16

Просто еще один способ делать декораторы. Я считаю, что это самый простой способ обернуть голову.

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()
Роберт Фей
источник
Спасибо! около 30 минут искал некоторые изумительные «решения», и это первое, что действительно имеет смысл.
canhazbits
0

Теперь, если вы хотите вызвать функцию function1с декоратором, decorator_with_argи в этом случае и функция, и декоратор принимают аргументы,

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

decorator_with_arg(10)(function1)(1, 2)
Сверхновая звезда
источник