Preservando assinaturas de funções decoradas

Suponha que eu tenha escrito um decorador que faz algo muito genérico. Por exemplo, ele pode converter todos os argumentos para um tipo específico, realizar registro, implementar memoização, etc.

Aqui está um exemplo:

def args_as_ints(f):
    def g(*args, **kwargs):
        args = [int(x) for x in args]
        kwargs = dict((k, int(v)) for k, v in kwargs.items())
        return f(*args, **kwargs)
    return g

@args_as_ints
def funny_function(x, y, z=3):
    """Computes x*y + 2*z"""
    return x*y + 2*z

>>> funny_function("3", 4.0, z="5")
22

Tudo bem até agora. No entanto, há um problema. A função decorada não retém a documentação da função original:

>>> help(funny_function)
Help on function g in module __main__:

g(*args, **kwargs)

Felizmente, existe uma solução alternativa:

def args_as_ints(f):
    def g(*args, **kwargs):
        args = [int(x) for x in args]
        kwargs = dict((k, int(v)) for k, v in kwargs.items())
        return f(*args, **kwargs)
    g.__name__ = f.__name__
    g.__doc__ = f.__doc__
    return g

@args_as_ints
def funny_function(x, y, z=3):
    """Computes x*y + 2*z"""
    return x*y + 2*z

Desta vez, o nome da função e a documentação estão corretos:

>>> help(funny_function)
Help on function funny_function in module __main__:

funny_function(*args, **kwargs)
    Computes x*y + 2*z

Mas ainda há um problema: a assinatura da função está errada. A informação "* args, ** kwargs" é quase inútil.

O que fazer? Posso pensar em duas soluções alternativas simples, mas falhas:

1 - Inclua a assinatura correta na docstring:

def funny_function(x, y, z=3):
    """funny_function(x, y, z=3) -- computes x*y + 2*z"""
    return x*y + 2*z

Isso é ruim por causa da duplicação. A assinatura ainda não será mostrada corretamente na documentação gerada automaticamente. É fácil atualizar a função e esquecer de alterar a docstring ou cometer um erro de digitação. [ E sim, estou ciente do fato de que a docstring já duplica o corpo da função. Ignore isso; funny_function é apenas um exemplo aleatório. ]

2 - Não use um decorador, ou use um decorador especial para cada assinatura específica:

def funny_functions_decorator(f):
    def g(x, y, z=3):
        return f(int(x), int(y), z=int(z))
    g.__name__ = f.__name__
    g.__doc__ = f.__doc__
    return g

Isso funciona bem para um conjunto de funções com assinatura idêntica, mas é inútil em geral. Como disse no início, quero poder usar decoradores de forma totalmente genérica.

Estou procurando uma solução totalmente geral e automática.

Portanto, a questão é: há uma maneira de editar a assinatura da função decorada depois que ela foi criada?

Caso contrário, posso escrever um decorador que extraia a assinatura da função e use essa informação em vez de "* kwargs, ** kwargs" ao construir a função decorada? Como faço para extrair essas informações? Como devo construir a função decorada - com exec?

Alguma outra abordagem?

1. Instale o módulo decorador :

   $ pip install decorator

2. Adapte a definição de args_as_ints():

   import decorator
   
   @decorator.decorator
   def args_as_ints(f, *args, **kwargs):
       args = [int(x) for x in args]
       kwargs = dict((k, int(v)) for k, v in kwargs.items())
       return f(*args, **kwargs)
   
   @args_as_ints
   def funny_function(x, y, z=3):
       """Computes x*y + 2*z"""
       return x*y + 2*z
   
   print funny_function("3", 4.0, z="5")
   # 22
   help(funny_function)
   # Help on function funny_function in module __main__:
   # 
   # funny_function(x, y, z=3)
   #     Computes x*y + 2*z

Python 3.4+

functools.wraps()de stdlib preserva assinaturas desde Python 3.4:

import functools


def args_as_ints(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        args = [int(x) for x in args]
        kwargs = dict((k, int(v)) for k, v in kwargs.items())
        return func(*args, **kwargs)
    return wrapper


@args_as_ints
def funny_function(x, y, z=3):
    """Computes x*y + 2*z"""
    return x*y + 2*z


print(funny_function("3", 4.0, z="5"))
# 22
help(funny_function)
# Help on function funny_function in module __main__:
#
# funny_function(x, y, z=3)
#     Computes x*y + 2*z

functools.wraps()está disponível pelo menos desde o Python 2.5, mas não preserva a assinatura lá:

help(funny_function)
# Help on function funny_function in module __main__:
#
# funny_function(*args, **kwargs)
#    Computes x*y + 2*z

Observe: em *args, **kwargsvez de x, y, z=3.

Isso é resolvido com a biblioteca padrão do Python functoolse, especificamente functools.wraps, a função, que é projetada para " atualizar uma função de wrapper para se parecer com a função de wrapper ". Seu comportamento depende da versão do Python, entretanto, conforme mostrado abaixo. Aplicado ao exemplo da pergunta, o código ficaria assim:

from functools import wraps

def args_as_ints(f):
    @wraps(f) 
    def g(*args, **kwargs):
        args = [int(x) for x in args]
        kwargs = dict((k, int(v)) for k, v in kwargs.items())
        return f(*args, **kwargs)
    return g


@args_as_ints
def funny_function(x, y, z=3):
    """Computes x*y + 2*z"""
    return x*y + 2*z

Quando executado em Python 3, isso produziria o seguinte:

>>> funny_function("3", 4.0, z="5")
22
>>> help(funny_function)
Help on function funny_function in module __main__:

funny_function(x, y, z=3)
    Computes x*y + 2*z

Sua única desvantagem é que no Python 2, no entanto, ele não atualiza a lista de argumentos da função. Quando executado em Python 2, produzirá:

>>> help(funny_function)
Help on function funny_function in module __main__:

funny_function(*args, **kwargs)
    Computes x*y + 2*z

Há um módulodecorator decorador com decorador que você pode usar:

@decorator
def args_as_ints(f, *args, **kwargs):
    args = [int(x) for x in args]
    kwargs = dict((k, int(v)) for k, v in kwargs.items())
    return f(*args, **kwargs)

Então, a assinatura e a ajuda do método são preservadas:

>>> help(funny_function)
Help on function funny_function in module __main__:

funny_function(x, y, z=3)
    Computes x*y + 2*z

EDIT: JF Sebastian apontou que eu não modifiquei a args_as_intsfunção - ela foi corrigida agora.

Dê uma olhada no módulo decorador - especificamente o decorador decorador, que resolve esse problema.

Segunda opçao:

1. Instale o módulo wrapt:

$ easy_install wrapt

wrapt tem um bônus, preserva a assinatura da classe.

import wrapt
import inspect


@wrapt.decorator
def args_as_ints(wrapped, instance, args, kwargs):
    if instance is None:
        if inspect.isclass(wrapped):
            # Decorator was applied to a class.
            return wrapped(*args, **kwargs)
        else:
            # Decorator was applied to a function or staticmethod.
            return wrapped(*args, **kwargs)
    else:
        if inspect.isclass(instance):
            # Decorator was applied to a classmethod.
            return wrapped(*args, **kwargs)
        else:
            # Decorator was applied to an instancemethod.
            return wrapped(*args, **kwargs)


@args_as_ints
def funny_function(x, y, z=3):
    """Computes x*y + 2*z"""
    return x * y + 2 * z


>>> funny_function(3, 4, z=5))
# 22

>>> help(funny_function)
Help on function funny_function in module __main__:

funny_function(x, y, z=3)
    Computes x*y + 2*z

Conforme comentado acima na resposta de jfs ; se você está preocupado com a assinatura em termos de aparência ( helpe inspect.signature), então o uso functools.wrapsé perfeitamente adequado.

Se você está preocupado com a assinatura em termos de comportamento (em particular TypeErrorno caso de incompatibilidade de argumentos), functools.wrapsnão a preserve. Você deve preferir usar decoratorpara isso, ou minha generalização de seu motor principal, chamado makefun.

from makefun import wraps

def args_as_ints(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        print("wrapper executes")
        args = [int(x) for x in args]
        kwargs = dict((k, int(v)) for k, v in kwargs.items())
        return func(*args, **kwargs)
    return wrapper


@args_as_ints
def funny_function(x, y, z=3):
    """Computes x*y + 2*z"""
    return x*y + 2*z


print(funny_function("3", 4.0, z="5"))
# wrapper executes
# 22

help(funny_function)
# Help on function funny_function in module __main__:
#
# funny_function(x, y, z=3)
#     Computes x*y + 2*z

funny_function(0)  
# observe: no "wrapper executes" is printed! (with functools it would)
# TypeError: funny_function() takes at least 2 arguments (1 given)

Veja também este post sobrefunctools.wraps .