Как создать Python-декоратор, который можно использовать с параметрами или без них?

Question

Я бы хотел создать декоратор Python, который можно использовать с параметрами:

@redirect_output("somewhere.log")
def foo():
    ....

или без них (например, для перенаправления вывода в stderr по умолчанию):

@redirect_output
def foo():
    ....

Это вообще возможно?

Обратите внимание, что я не ищу другого решения проблемы перенаправления вывода, это всего лишь пример синтаксиса, которого я хотел бы достичь.

Answer 1

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

Как и в ответе Тобе, единственный способ справиться с обоими случаями - проверить оба сценария. Самый простой способ - просто проверить, существует ли один аргумент и является ли он callabe (ПРИМЕЧАНИЕ: дополнительные проверки будут необходимы, если ваш декоратор принимает только 1 аргумент и это может быть вызываемый объект):

def decorator(*args, **kwargs):
    if len(args) == 1 and len(kwargs) == 0 and callable(args[0]):
        # called as @decorator
    else:
        # called as @decorator(*args, **kwargs)

В первом случае вы делаете то, что делает любой обычный декоратор, возвращаете модифицированную или упакованную версию переданной функции.

Во втором случае вы возвращаете «новый» декоратор, который каким-то образом использует информацию, переданную с помощью * args, ** kwargs.

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

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

def doublewrap(f):
    '''
    a decorator decorator, allowing the decorator to be used as:
    @decorator(with, arguments, and=kwargs)
    or
    @decorator
    '''
    @wraps(f)
    def new_dec(*args, **kwargs):
        if len(args) == 1 and len(kwargs) == 0 and callable(args[0]):
            # actual decorated function
            return f(args[0])
        else:
            # decorator arguments
            return lambda realf: f(realf, *args, **kwargs)

    return new_dec

Теперь мы можем украсить наши декораторы с помощью @doublewrap, и они будут работать с аргументами и без них, с одной оговоркой:

Я отметил выше, но должен повторить здесь, проверка в этом декораторе делает предположение об аргументах, которые может получить декоратор (а именно, что он не может получить один вызываемый аргумент). Поскольку мы делаем его применимым к любому генератору сейчас, его нужно помнить или модифицировать, если это будет противоречить.

Следующее демонстрирует его использование:

def test_doublewrap():
    from util import doublewrap
    from functools import wraps    

    @doublewrap
    def mult(f, factor=2):
        '''multiply a function's return value'''
        @wraps(f)
        def wrap(*args, **kwargs):
            return factor*f(*args,**kwargs)
        return wrap

    # try normal
    @mult
    def f(x, y):
        return x + y

    # try args
    @mult(3)
    def f2(x, y):
        return x*y

    # try kwargs
    @mult(factor=5)
    def f3(x, y):
        return x - y

    assert f(2,3) == 10
    assert f2(2,5) == 30
    assert f3(8,1) == 5*7

Answer 2

Использование аргументов ключевого слова со значениями по умолчанию (как предложено kquinn) - хорошая идея, но вам потребуется включить скобку:

@redirect_output()
def foo():
    ...

Если вам нужна версия, которая работает без скобок в декораторе, вам придется учитывать оба сценария в коде декоратора.

Если бы вы использовали Python 3.0, вы могли бы использовать аргументы только для ключевых слов:

def redirect_output(fn=None,*,destination=None):
  destination = sys.stderr if destination is None else destination
  def wrapper(*args, **kwargs):
    ... # your code here
  if fn is None:
    def decorator(fn):
      return functools.update_wrapper(wrapper, fn)
    return decorator
  else:
    return functools.update_wrapper(wrapper, fn)

В Python 2.x это можно эмулировать с помощью трюков varargs:

def redirected_output(*fn,**options):
  destination = options.pop('destination', sys.stderr)
  if options:
    raise TypeError("unsupported keyword arguments: %s" % 
                    ",".join(options.keys()))
  def wrapper(*args, **kwargs):
    ... # your code here
  if fn:
    return functools.update_wrapper(wrapper, fn[0])
  else:
    def decorator(fn):
      return functools.update_wrapper(wrapper, fn)
    return decorator

Любая из этих версий позволит вам написать код, подобный этому:

@redirected_output
def foo():
    ...

@redirected_output(destination="somewhere.log")
def bar():
    ...

Answer 3

Вам необходимо обнаружить оба случая, например, используя тип первого аргумента, и, соответственно, вернуть либо оболочку (при использовании без параметров), либо декоратор (при использовании с аргументами).

from functools import wraps
import inspect

def redirect_output(fn_or_output):
    def decorator(fn):
        @wraps(fn)
        def wrapper(*args, **args):
            # Redirect output
            try:
                return fn(*args, **args)
            finally:
                # Restore output
        return wrapper

    if inspect.isfunction(fn_or_output):
        # Called with no parameter
        return decorator(fn_or_output)
    else:
        # Called with a parameter
        return decorator

При использовании синтаксиса @redirect_output("output.log"), redirect_output вызывается с одним аргументом "output.log", и он должен возвращать декоратор, принимающий функцию, которая будет оформлена в качестве аргумента. Когда используется как @redirect_output, он вызывается напрямую с функцией, которая должна быть оформлена как аргумент.

Или другими словами: за синтаксисом @ должно следовать выражение, результатом которого является функция, принимающая функцию, которая должна быть декорирована в качестве единственного аргумента, и возвращающая декорированную функцию. Само выражение может быть вызовом функции, как в случае @redirect_output("output.log"). Запутанно, но верно: -)

Answer 4

Декоратор Python вызывается принципиально по-разному, в зависимости от того, даете ли вы ему аргументы или нет. Декорация - это просто (синтаксически ограниченное) выражение.

В вашем первом примере:

@redirect_output("somewhere.log")
def foo():
    ....

функция redirect_output вызывается с заданный аргумент, который должен вернуть декоратор функция, которая сама вызывается с foo в качестве аргумента, которая (наконец-то!) должна вернуть окончательно оформленную функцию.

Эквивалентный код выглядит так:

def foo():
    ....
d = redirect_output("somewhere.log")
foo = d(foo)

Эквивалентный код для вашего второго примера выглядит так:

def foo():
    ....
d = redirect_output
foo = d(foo)

Таким образом, вы можете делать то, что вам нравится, но не совсем гладко:

import types
def redirect_output(arg):
    def decorator(file, f):
        def df(*args, **kwargs):
            print 'redirecting to ', file
            return f(*args, **kwargs)
        return df
    if type(arg) is types.FunctionType:
        return decorator(sys.stderr, arg)
    return lambda f: decorator(arg, f)

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

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

@redirect_output()
def foo():
    ....

Код функции декоратора будет выглядеть так:

def redirect_output(file = sys.stderr):
    def decorator(file, f):
        def df(*args, **kwargs):
            print 'redirecting to ', file
            return f(*args, **kwargs)
        return df
    return lambda f: decorator(file, f)

Answer 5

Я знаю, что это старый вопрос, но мне действительно не нравится ни один из предложенных методов, поэтому я хотел добавить другой метод. Я видел, что django использует действительно чистый метод в их login_required декораторе в django.contrib.auth.decorators. Как вы можете видеть в документах декоратора , его можно использовать отдельно как @login_required или с аргументами @login_required(redirect_field_name='my_redirect_field').

То, как они это делают, довольно просто. Они добавляют kwarg (function=None) перед аргументами декоратора. Если декоратор используется один, function будет фактической функцией, которую он декорирует, тогда как если он вызывается с аргументами, function будет None.

Пример:

from functools import wraps

def custom_decorator(function=None, some_arg=None, some_other_arg=None):
    def actual_decorator(f):
        @wraps(f)
        def wrapper(*args, **kwargs):
            # Do stuff with args here...
            if some_arg:
                print(some_arg)
            if some_other_arg:
                print(some_other_arg)
            return f(*args, **kwargs)
        return wrapper
    if function:
        return actual_decorator(function)
    return actual_decorator

---

@custom_decorator
def test1():
    print('test1')

>>> test1()
test1

---

@custom_decorator(some_arg='hello')
def test2():
    print('test2')

>>> test2()
hello
test2

---

@custom_decorator(some_arg='hello', some_other_arg='world')
def test3():
    print('test3')

>>> test3()
hello
world
test3

Я считаю, что такой подход, который использует django, более элегантен и прост для понимания, чем любой другой метод, предложенный здесь.

Answer 6

Несколько ответов здесь уже хорошо решают вашу проблему. Что касается стиля, однако, я предпочитаю решить эту проблему с использованием декоратора functools.partial, как предложено в Python Cookbook 3 Дэвида Бизли:

from functools import partial, wraps

def decorator(func=None, foo='spam'):
    if func is None:
         return partial(decorator, foo=foo)

    @wraps(func)
    def wrapper(*args, **kwargs):
        # do something with `func` and `foo`, if you're so inclined
        pass

    return wrapper

Хотя да, вы можете просто сделать

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

без прикольных обходных путей, я нахожу это странно выглядящим, и мне нравится иметь возможность просто декорировать @decorator.

Что касается вторичной цели миссии, перенаправление вывода функции рассматривается в этом Пост переполнения стека .

---

Если вы хотите погрузиться глубже, ознакомьтесь с главой 9 (Метапрограммирование) в Python Cookbook 3 , которая свободно доступна для чтения онлайн .

Некоторые из этих материалов демонстрируются вживую (и еще больше!) В потрясающем видео Beazley на YouTube Метапрограммирование Python 3 .

Удачного кодирования:)

Answer 7

Фактически, случай предостережения в решении @ bj0 можно легко проверить:

def meta_wrap(decor):
    @functools.wraps(decor)
    def new_decor(*args, **kwargs):
        if len(args) == 1 and len(kwargs) == 0 and callable(args[0]):
            # this is the double-decorated f. 
            # Its first argument should not be a callable
            doubled_f = decor(args[0])
            @functools.wraps(doubled_f)
            def checked_doubled_f(*f_args, **f_kwargs):
                if callable(f_args[0]):
                    raise ValueError('meta_wrap failure: '
                                'first positional argument cannot be callable.')
                return doubled_f(*f_args, **f_kwargs)
            return checked_doubled_f 
        else:
            # decorator arguments
            return lambda real_f: decor(real_f, *args, **kwargs)

    return new_decor

Вот несколько тестов для этой отказоустойчивой версии meta_wrap.

    @meta_wrap
    def baddecor(f, caller=lambda x: -1*x):
        @functools.wraps(f)
        def _f(*args, **kwargs):
            return caller(f(args[0]))
        return _f

    @baddecor  # used without arg: no problem
    def f_call1(x):
        return x + 1
    assert f_call1(5) == -6

    @baddecor(lambda x : 2*x) # bad case
    def f_call2(x):
        return x + 1
    f_call2(5)  # raises ValueError

    # explicit keyword: no problem
    @baddecor(caller=lambda x : 100*x)
    def f_call3(x):
        return x + 1
    assert f_call3(5) == 600

Answer 8

Опираясь на ответ Вартека:

imports sys

def redirect_output(func, output=None):
    if output is None:
        output = sys.stderr
    if isinstance(output, basestring):
        output = open(output, 'w') # etc...
    # everything else...

Answer 9

Вы пробовали аргументы ключевых слов со значениями по умолчанию? Что-то вроде

def decorate_something(foo=bar, baz=quux):
    pass

Answer 10

Обычно вы можете задать аргументы по умолчанию в Python ...

def redirect_output(fn, output = stderr):
    # whatever

Не уверен, что это работает и с декораторами. Я не знаю ни одной причины, почему это не так.