Como criar um decorador Python que pode ser usado com ou sem parâmetros?

Eu gostaria de criar um decorador Python que pode ser usado com parâmetros:

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

ou sem eles (por exemplo, para redirecionar a saída para stderr por padrão):

@redirect_output
def foo():
    ....

isso é, de todo, possível?

Observe que não estou procurando uma solução diferente para o problema de redirecionamento de saída, é apenas um exemplo da sintaxe que gostaria de obter.

Sei que essa pergunta é antiga, mas alguns dos comentários são novos e, embora todas as soluções viáveis ​​sejam essencialmente as mesmas, a maioria delas não é muito clara ou fácil de ler.

Como diz a resposta de Thobe, a única maneira de lidar com os dois casos é verificar os dois cenários. A maneira mais fácil é simplesmente verificar se há um único argumento e se ele pode ser chamado (OBSERVAÇÃO: verificações extras serão necessárias se o seu decorador receber apenas 1 argumento e for um objeto que pode ser chamado):

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)

No primeiro caso, você faz o que qualquer decorador normal faz, retorna uma versão modificada ou encapsulada da função passada.

No segundo caso, você retorna um 'novo' decorador que de alguma forma usa as informações passadas com * args, ** kwargs.

Isso é bom e tudo, mas ter que escrever isso para cada decorador que você fizer pode ser muito chato e não tão limpo. Em vez disso, seria bom ser capaz de modificar automaticamente nossos decoradores sem ter que reescrevê-los ... mas é para isso que servem os decoradores!

Usando o seguinte decorador decorador, podemos deocrate nossos decoradores para que eles possam ser usados ​​com ou sem argumentos:

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

Agora, podemos decorar nossos decoradores com @doublewrap, e eles trabalharão com e sem argumentos, com uma ressalva:

Eu observei acima, mas devo repetir aqui, a verificação neste decorador faz uma suposição sobre os argumentos que um decorador pode receber (ou seja, que ele não pode receber um único argumento que pode ser chamado). Já que o estamos tornando aplicável a qualquer gerador agora, ele precisa ser mantido em mente ou modificado se for contestado.

O seguinte demonstra seu uso:

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

Usar argumentos de palavra-chave com valores padrão (como sugerido por kquinn) é uma boa ideia, mas exigirá que você inclua o parêntese:

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

Se desejar uma versão que funcione sem parênteses no decorador, você terá que levar em conta os dois cenários em seu código de decorador.

Se você estivesse usando Python 3.0, poderia usar argumentos apenas de palavra-chave para isso:

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)

No Python 2.x, isso pode ser emulado com truques de 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

Qualquer uma dessas versões permitiria que você escrevesse um código como este:

@redirected_output
def foo():
    ...

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

Eu sei que esta é uma questão antiga, mas eu realmente não gosto de nenhuma das técnicas propostas, então eu queria adicionar outro método. Eu vi que o django usa um método realmente limpo em seu login_requireddecorador emdjango.contrib.auth.decorators . Como você pode ver nos documentos do decorador , ele pode ser usado sozinho @login_requiredou com argumentos @login_required(redirect_field_name='my_redirect_field'),.

A forma como o fazem é bastante simples. Eles adicionam um kwarg( function=None) antes de seus argumentos de decorador. Se o decorador for usado sozinho, functionserá a função real que ele está decorando, enquanto se for chamado com argumentos, functionserá None.

Exemplo:

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

Acho que essa abordagem que django usa é mais elegante e fácil de entender do que qualquer uma das outras técnicas propostas aqui.

Você precisa detectar ambos os casos, por exemplo, usando o tipo do primeiro argumento e, portanto, retornar o invólucro (quando usado sem parâmetro) ou um decorador (quando usado com argumentos).

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

Ao usar a @redirect_output("output.log")sintaxe, redirect_outputé chamado com um único argumento "output.log", e deve retornar um decorador aceitando a função a ser decorada como argumento. Quando usado como @redirect_output, é chamado diretamente com a função a ser decorada como um argumento.

Ou seja: a @sintaxe deve ser seguida por uma expressão cujo resultado seja uma função que aceita uma função a ser decorada como seu único argumento e retorna a função decorada. A própria expressão pode ser uma chamada de função, que é o caso de @redirect_output("output.log"). Convoluto, mas verdadeiro :-)

Várias respostas aqui já resolvem seu problema de maneira adequada. No que diz respeito ao estilo, no entanto, prefiro resolver esse problema do decorador usando functools.partial, conforme sugerido no Python Cookbook 3 de David Beazley :

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

Embora sim, você pode apenas fazer

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

sem soluções alternativas, acho estranho e gosto de ter a opção de simplesmente decorar com @decorator.

Quanto ao objetivo da missão secundária, redirecionar a saída de uma função é abordado nesta postagem do Stack Overflow .

Se você quiser se aprofundar, verifique o Capítulo 9 (Metaprogramação) no Python Cookbook 3 , que está disponível gratuitamente para leitura online .

Parte desse material é demo ao vivo (e mais!) No incrível vídeo de Beazley no YouTube em Python 3 Metaprogramming .

Boa codificação :)

Um decorador python é chamado de uma maneira fundamentalmente diferente, dependendo de você fornecer argumentos ou não. A decoração é na verdade apenas uma expressão (restrita sintaticamente).

Em seu primeiro exemplo:

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

a função redirect_outputé chamada com o argumento fornecido, que deve retornar uma função decoradora, que por sua vez é chamada com fooum argumento, que (finalmente!) deve retornar a função decorada final.

O código equivalente se parece com este:

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

O código equivalente para seu segundo exemplo se parece com:

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

Assim, você pode fazer o que quiser, mas não de uma forma totalmente contínua:

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)

Isso deve funcionar, a menos que você deseje usar uma função como um argumento para o seu decorador, caso em que o decorador presumirá erroneamente que ela não tem argumentos. Também falhará se esta decoração for aplicada a outra decoração que não retorna um tipo de função.

Um método alternativo é apenas exigir que a função decoradora seja sempre chamada, mesmo que não seja nenhum argumento. Nesse caso, seu segundo exemplo seria assim:

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

O código da função do decorador seria assim:

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)

Na verdade, o caso de advertência na solução de @ bj0 pode ser verificado facilmente:

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

Aqui estão alguns casos de teste para esta versão à prova de falhas do 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

Para dar uma resposta mais completa do que a anterior:

"Existe uma maneira de construir um decorador que pode ser usado com e sem argumentos?"

Não, não há uma maneira genérica porque atualmente falta algo na linguagem python para detectar os dois casos de uso diferentes.

No entanto, Sim, conforme já apontado por outras respostas como bj0s , há uma solução alternativa desajeitada que é verificar o tipo e o valor do primeiro argumento posicional recebido (e verificar se nenhum outro argumento tem valor não padrão). Se você tiver a garantia de que os usuários nunca passarão um chamável como primeiro argumento do decorador, você pode usar esta solução alternativa. Observe que isso é o mesmo para decoradores de classe (substitua chamável por classe na frase anterior).

Para ter certeza do que foi dito acima, eu pesquisei bastante e até implementei uma biblioteca chamada decopatchque usa uma combinação de todas as estratégias citadas acima (e muito mais, incluindo a introspecção) para realizar "qualquer que seja a solução alternativa mais inteligente", dependendo em sua necessidade.

Mas, francamente, o melhor seria não precisar de nenhuma biblioteca aqui e obter esse recurso diretamente da linguagem python. Se, como eu, você acha que é uma pena que a linguagem python não seja, até hoje, capaz de fornecer uma resposta clara a essa pergunta, não hesite em apoiar essa ideia no bugtracker do python : https: //bugs.python .org / issue36553 !

Muito obrigado por sua ajuda para tornar o python uma linguagem melhor :)

Isso faz o trabalho sem complicações:

from functools import wraps

def memoize(fn=None, hours=48.0):
  def deco(fn):
    @wraps(fn)
    def wrapper(*args, **kwargs):
      return fn(*args, **kwargs)
    return wrapper

  if callable(fn): return deco(fn)
  return deco

Como ninguém mencionou isso, há também uma solução utilizando a classe chamável que considero mais elegante, especialmente nos casos em que o decorador é complexo e pode-se querer dividi-lo em vários métodos (funções). Esta solução utiliza um __new__método mágico para fazer essencialmente o que os outros indicaram. Primeiro, detecte como o decorador foi usado e depois ajuste o retorno de forma adequada.

class decorator_with_arguments(object):

    def __new__(cls, decorated_function=None, **kwargs):

        self = super().__new__(cls)
        self._init(**kwargs)

        if not decorated_function:
            return self
        else:
            return self.__call__(decorated_function)

    def _init(self, arg1="default", arg2="default", arg3="default"):
        self.arg1 = arg1
        self.arg2 = arg2
        self.arg3 = arg3

    def __call__(self, decorated_function):

        def wrapped_f(*args):
            print("Decorator arguments:", self.arg1, self.arg2, self.arg3)
            print("decorated_function arguments:", *args)
            decorated_function(*args)

        return wrapped_f

@decorator_with_arguments(arg1=5)
def sayHello(a1, a2, a3, a4):
    print('sayHello arguments:', a1, a2, a3, a4)

@decorator_with_arguments()
def sayHello(a1, a2, a3, a4):
    print('sayHello arguments:', a1, a2, a3, a4)

@decorator_with_arguments
def sayHello(a1, a2, a3, a4):
    print('sayHello arguments:', a1, a2, a3, a4)

Se o decorador for usado com argumentos, isso será igual a:

result = decorator_with_arguments(arg1=5)(sayHello)(a1, a2, a3, a4)

Pode-se ver que os argumentos arg1são passados ​​corretamente para o construtor e a função decorada é passada para__call__

Mas se o decorador for usado sem argumentos, isso será igual a:

result = decorator_with_arguments(sayHello)(a1, a2, a3, a4)

Você vê que, neste caso, a função decorada é passada diretamente para o construtor e a chamada para __call__é totalmente omitida. É por isso que precisamos empregar a lógica para cuidar desse caso no __new__método mágico.

Por que não podemos usar em __init__vez de __new__? O motivo é simples: o python proíbe o retorno de qualquer outro valor além de None de__init__

ATENÇÃO

Essa abordagem tem um efeito colateral. Não preservará a assinatura da função!

Você tentou argumentos de palavra-chave com valores padrão? Algo como

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

Geralmente você pode fornecer argumentos padrão em Python ...

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

Não tenho certeza se isso funciona bem com decoradores, no entanto. Não sei por que não.

Com base na resposta da vartec:

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...