在装饰器中修改函数

Posted 2023-02-23

【中文标题】在装饰器中修改函数

【英文标题】：Modify function in decorator

【发布时间】：2015-09-13 17:14:07

【问题描述】：

我正在考虑制作一个装饰器以提高性能。一个装饰器，它修改它所装饰的函数的源代码，并返回修改后的函数。

在思考这个问题时，我想如果我能得到函数的源代码，我就可以做到这一点。但是是否可以在装饰器中访问函数的源代码？如果我有这样的装饰器：

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

【问题讨论】：

ast 可能更适合你想要的东西

在exec(inspect.getsource(f)) 之前留一行空，它会起作用。有趣！！

@AshwiniChaudhary 这对我没有任何作用......但这让我认为这可能是一个并发问题，或者依赖于实现

@Quincunx 我在 IPython shell 中使用 %run 进行了不同的测试。使用普通的 python 命令它不起作用。

@Quincunx，你打算如何修改源代码？

【参考方案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 库，有一个很棒的 video from pycon 2011，其中 Matthew Desmarais 讲述了如何使用 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 失败，因为它找不到定义在字符串。

【讨论】：

啊，有道理

这很聪明。看起来exec 正在查看装饰器并递归应用它，但失败了。