¿Cómo crear un decorador de Python que se pueda usar con o sin parámetros?

Question

Me gustaría crear un decorador de Python que se pueda usar con parámetros:

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

o sin ellos (por ejemplo, para redirigir la salida a stderr de forma predeterminada):

@redirect_output
def foo():
    ....

¿Es eso posible?

Tenga en cuenta que no estoy buscando una solución diferente al problema de redirigir la salida, es solo un ejemplo de la sintaxis que me gustaría lograr.

Answer 1

Sé que esta pregunta es antigua, pero algunos de los comentarios son nuevos y, si bien todas las soluciones viables son esencialmente las mismas, la mayoría de ellas no son muy limpias o fáciles de leer.

Como dice la respuesta de thobe, la única forma de manejar ambos casos es verificar ambos escenarios. La forma más fácil es simplemente verificar si hay un argumento único y es callabe (NOTA: se necesitarán verificaciones adicionales si su decorador solo toma 1 argumento y resulta ser un objeto invocable):

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)

En el primer caso, hace lo que hace cualquier decorador normal, devuelve una versión modificada o envuelta de la función pasada.

En el segundo caso, devuelve un 'nuevo' decorador que de alguna manera usa la información que se pasa con * args, ** kwargs.

Esto está bien, pero tener que escribirlo para cada decorador que hagas puede ser bastante molesto y no tan limpio. En cambio, sería bueno poder modificar automáticamente nuestros decoradores sin tener que volver a escribirlos ... ¡pero para eso están los decoradores!

Usando el siguiente decorador decorador, podemos desocular nuestros decoradores para que puedan usarse con o sin 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

Ahora, podemos decorar nuestros decoradores con @doublewrap, y funcionarán con y sin argumentos, con una advertencia:

Observé anteriormente, pero debería repetir aquí, la comprobación en este decorador asume los argumentos que un decorador puede recibir (es decir, que no puede recibir un único argumento invocable). Dado que ahora lo estamos aplicando a cualquier generador, debe tenerse en cuenta o modificarse si se contradice.

Lo siguiente demuestra su 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

Answer 2

Usar argumentos de palabras clave con valores predeterminados (como sugiere kquinn) es una buena idea, pero requerirá que incluya el paréntesis:

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

Si desea una versión que funcione sin paréntesis en el decorador, deberá tener en cuenta ambos escenarios en su código de decorador.

Si estaba usando Python 3.0, podría usar argumentos de solo palabras clave para esto:

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)

En Python 2.x esto se puede emular con trucos 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

Cualquiera de estas versiones le permitiría escribir código como este:

@redirected_output
def foo():
    ...

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

Answer 3

Debe detectar ambos casos, por ejemplo, utilizando el tipo del primer argumento y, en consecuencia, devolver el contenedor (cuando se usa sin parámetro) o un decorador (cuando se usa con 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

Cuando se utiliza la sintaxis @redirect_output (" output.log ") , se invoca redirect_output con un solo argumento " output.log " , y debe devolver un decorador que acepte la función que se decorará como argumento. Cuando se usa como @redirect_output , se llama directamente con la función que se decorará como argumento.

O en otras palabras: la sintaxis @ debe ir seguida de una expresión cuyo resultado sea una función que acepte una función para decorar como único argumento y que devuelva la función decorada. La expresión en sí puede ser una llamada de función, como es el caso de @redirect_output (" output.log ") . Enrevesado, pero cierto :-)

Answer 4

Un decorador de python se llama de una manera fundamentalmente diferente dependiendo de si le das argumentos o no. La decoración es en realidad solo una expresión (sintácticamente restringida).

En su primer ejemplo:

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

la función redirect_output se llama con el argumento dado, que se espera que devuelva un decorador función, que se llama con foo como argumento, que (¡finalmente!) se espera que devuelva la función decorada final.

El código equivalente se ve así:

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

El código equivalente para su segundo ejemplo se ve así:

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

Entonces puede hacer lo que quiera pero no de manera totalmente fluida:

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)

Esto debería estar bien a menos que desee utilizar una función como argumento a su decorador, en cuyo caso el decorador asumirá erróneamente que no tiene argumentos. También fallará si esta decoración se aplica a otra decoración que no devuelve un tipo de función.

Un método alternativo es simplemente requerir que el La función decoradora siempre se llama, incluso si no tiene argumentos. En este caso, su segundo ejemplo se vería así:

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

El código de la función del decorador se vería así:

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

Sé que esta es una vieja pregunta, pero realmente no me gusta ninguna de las técnicas propuestas, así que quería agregar otro método. Vi que django usa un método realmente limpio en su login_required en django.contrib.auth.decorators . Como puede ver en los documentos del decorador , se puede usar solo como @login_required o con argumentos, @login_required (redirect_field_name = 'my_redirect_field') .

La forma en que lo hacen es bastante simple. Agregan un kwarg ( function = None ) antes de sus argumentos de decorador. Si el decorador se usa solo, function será la función real que está decorando, mientras que si se llama con argumentos, function será None .

Ejemplo :

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

Encuentro que este enfoque que usa django es más elegante y fácil de entender que cualquiera de las otras técnicas propuestas aquí.

Answer 6

Varias respuestas aquí ya abordan su problema muy bien. Sin embargo, con respecto al estilo, prefiero resolver este problema de decorador usando functools.partial , como se sugiere en el 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

Mientras que sí, solo puedes hacer

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

sin soluciones funky, me parece extraño, y me gusta tener la opción de simplemente decorar con @decorator .

En cuanto al objetivo de la misión secundaria, la redirección de la salida de una función se aborda en este Publicación de desbordamiento de pila .

---

Si desea profundizar, consulte el Capítulo 9 (Metaprogramación) en Python Cookbook 3 , que está disponible gratuitamente para ser leer en línea .

Parte de ese material se muestra en vivo (¡y más!) en el asombroso video de YouTube de Beazley Metaprogramación de Python 3 .

Feliz codificación :)

Answer 7

De hecho, el caso de advertencia en la solución de @ bj0 se puede verificar fácilmente:

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

Aquí hay algunos casos de prueba para esta versión a prueba de fallas de 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

¿Has probado argumentos de palabras clave con valores predeterminados? Algo como

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

Answer 9

Generalmente puede dar argumentos predeterminados en Python ...

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

No estoy seguro si eso también funciona con decoradores. No sé de ninguna razón por la que no lo haría.

Answer 10

Basándose en la respuesta de 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...