Изменить функцию в декораторе

Я думал о создании декоратора с целью повышения производительности. Декоратор, который изменяет исходный код функции, которую он украшает, и возвращает измененную функцию.

Подумав об этом, я подумал, что если бы я мог просто получить исходный код функции, я мог бы это сделать. Но возможно ли получить доступ к исходному коду функции внутри декоратора? Если у меня есть декоратор:

import inspect

def decorate(f):
    exec(inspect.getsource(f))
    return eval(f.__name__)

@decorate
def test():
    return 1

Я получаю OSError:

OSError: could not get source code

Это похоже на то, что test не сформировалось полностью до того, как оно будет передано в decorate. Однако это работает:

import inspect

def decorate(f):
    exec(inspect.getsource(f))
    return eval(f.__name__)

def test():
    return 1
test = decorate(test)

Просто у него нет этого декоратора. Кажется, что это возможно, потому что f.__code__ определено.

При дальнейшем осмотре кажется, что это происходит только тогда, когда я помещаю inspect.getsource(f) в exec. В противном случае кажется, что я могу получить исходный код.

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

def tail_recurse(acc_default):
    def decorate(f):
        def wrapper(*args, acc=acc_default):
            args = args + (acc,)
            while True:
                return_type, *rargs = f(*args)
                if return_type is None:
                    return rargs[-1]
                args = rargs
        return wrapper
    return decorate

В принципе, я собираюсь сделать что-то так же просто, как заменить тело функции:

while True:
    __body__
    update_args

Ответ 1

Вы можете использовать functools.wraps с исходным кодом:

import inspect
from functools import wraps

@wraps
def decorate(f):
    exec(inspect.getsource(f))
    return eval(f.__name__)

@decorate
def test():
    return 1

Вывод:

In [2]: test()
Out[2]: 1

Если вы планируете менять источник во время выполнения, вы должны ознакомиться с библиотекой ast, есть отличный видео с pycon 2011, где Мэтью Десмараис рассказывает о том, как использовать модуль ast для изменения исходного кода с основ, вплоть до более расширенных параметров, это это простой рабочий пример используемого в разговоре python для javascript-переводчика, он будет работать для простых примеров, таких как функция fib.

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

from ast import parse, NodeTransformer


class Transformer(NodeTransformer):
    def __init__(self):
        self.src = ""
        self.indent = 0

    def translate(self, node):
        self.visit(node)
        return self.src

    def _indent(self, line):
        return "{}{line}".format(" " * self.indent, line=line)

    def render(self, body):
        self.indent += 2
        for stmt in body:
            self.visit(stmt)
        self.indent -= 2

    def visit_Num(self, node):
        self.src += "{}".format(node.n)

    def visit_Str(self, node):
        self.src += "{}".format(node.s)

    def visit_FunctionDef(self, defn):
        args = ",".join(name.arg for name in defn.args.args)
        js_defn = "var {} = function({}){{\n"
        self.src += self._indent(js_defn.format(defn.name, args))
        self.render(defn.body)
        self.src += self._indent("}\n")

    def visit_Eq(self, less):
        self.src += "=="

    def visit_Name(self, name):
        self.src += "{}".format(name.id)

    def visit_BinOp(self, binop):
        self.visit(binop.left)
        self.src += " "
        self.visit(binop.op)
        self.src += " "
        self.visit(binop.right)

    def visit_If(self, _if):
        self.src += self._indent("if (")
        self.visit(_if.test)
        self.src += ") {\n"
        self.render(_if.body)
           self.src += " "*self.indent + "}\n"


    def visit_Compare(self, comp):
        self.visit(comp.left)
        self.src += " "
        self.visit(comp.ops[0])
        self.src += " "
        self.visit(comp.comparators[0])

    def visit_Call(self, call):
        self.src += " "
        self.src += "{}(".format(call.func.id)
        self.visit(call.args[0])
        self.src += ")"

    def visit_Add(self, add):
        self.src += "+"

    def visit_Sub(self, add):
        self.src += "-"

    def visit_Return(self, ret):
        self.src += self._indent("return")
        if ret.value:
            self.src += " "
            self.visit(ret.value)
        self.src += ";\n"


def dec(f):
    source = getsource(f)
    _ast = parse(source)
    trans = Transformer()
    trans.indent = 0
    return trans.translate(_ast)


from inspect import getsource


def fibonacci(n):
    if n == 0:
        return 0
    if n == 1:
        return 1
    return fibonacci(n - 1) + fibonacci(n - 2)

Запуск функции dec выводит наш python как javascript:

print(dec(fibonacci))
var fibonacci = function(n){
  if (n == 0) {
    return 0;
  }
  if (n == 1) {
    return 1;
  }
  return  fibonacci(n - 1) +  fibonacci(n - 2);
}

Документы greentreesnakes также заслуживают внимания.

Ответ 2

Это работает:

import inspect, itertools

def decorate(f):
    source = itertools.dropwhile(lambda line: line.startswith('@'), inspect.getsource(f).splitlines())
    exec('\n'.join(source))
    return eval(f.__name__)

@decorate
def test():
    return 1

Я думаю, что проблема заключается в включении декоратора в источник функции.

# foo.py
import inspect

def decorate(f):
    print inspect.getsource(f)

@decorate
def test():
    return 1

>>> import foo
@decorate
def test():
    return 1
>>> # Notice the decorator is included in the source.

exec видит @decorate для test, определенного в строке, поэтому он рекурсивно вызывает decorate, но на этот раз inspect.getsource выходит из строя, поскольку он не может найти источник функции, определенной в строка.