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

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

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

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__ is определенными.

при дальнейшем рассмотрении оказывается, что это происходит только тогда, когда я ставлю 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
10
decorator function python python-decorators

2 ответов

можно использовать functools.обертывания С вашим исходным кодом:

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

если вы планируете изменить источник во время выполнения, то вы должны ознакомиться с АСТ библиотека, есть отличный видео из команда 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 документы также стоит прочитать.

4

это работает:

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 не удается, потому что он не может найти источник функции, определенной в строке.

1