Python通过函数装饰器转换ast

Question

我有一个想法，可以通过类似下面的装饰器来转换所有给定的使用装饰器标记的函数:

@transform_ast
def foo(x):
    return x

在 transform_ast 中，我获取源代码，提取抽象语法树（AST），对其进行变换，然后再次从中创建代码对象和函数类型。大致如下所示：

import ast
import inspect
import types

class Rewrite(ast.NodeTransformer):
    pass

def transform_ast(f):

    source = inspect.getsource(f)
    source = '\n'.join(source.splitlines()[1:]) # remove the decorator first line.
    print(source)

    old_code_obj = f.__code__
    old_ast = ast.parse(source)
    new_ast = Rewrite().visit(old_ast)
    new_code_obj = compile(new_ast, old_code_obj.co_filename, 'exec')
    new_f = types.FunctionType(new_code_obj, {})
    return new_f

@transform_ast
def foo(x):
    return x

然而，当我随后调用foo(x)时，它似乎不能正常工作。

实际上，我们可以假设我的转换只是将return x重写为return x+1。理想情况下，我希望一切都像往常一样正常工作，包括能够使用调试器进入函数...

调用foo(10)会出现以下错误：

TypeError: module() takes no arguments (1 given)

我做错了什么吗？

- Danny

Answer 1

new_code_obj = compile(new_ast, old_code_obj.co_filename, 'exec')

使用exec模式编译的代码始终被视为模块级别代码，当然它可以包含函数或类定义，或任何其他有效的Python代码。

要验证这一点，您可以访问代码对象的co_name属性，以获取定义此代码对象的名称。

>>> new_code_obj.co_name
<module>

在这里，new_code_obj是对应模块的代码对象。但是函数foo的代码对象在哪里呢？我们如何访问它？

可以从代码对象的co_consts属性中访问，该属性是一个用于字节码的常量元组。

>>> new_code_obj.co_consts
(<code object foo at 0x031C3DE0, file "c:/Users/test.py", line 1>, 'foo', None)
>>> new_code_obj.co_consts[0]
<code object foo at 0x031C3DE0, file "c:/Users/test.py", line 1>

要验证这个代码对象来自函数foo，你可以再次使用co_name属性。

>>> new_code_obj.co_consts[0].co_name
foo

因此，在创建新的FunctionType时，您应该使用与函数foo对应的代码对象，而不是module代码对象。

这样做会改变

new_f = types.FunctionType(new_code_obj, {})

为了

new_f = types.FunctionType(new_code_obj.co_consts[0], f.__globals__)
# Here `f` is the function object passed to the `transform_ast`

我将解决这个问题。

其他参考资料：探索Python代码对象

- Abdul Niyas P M

---

这样不行，尝试使用装饰器与 def foo(): return range(10) 一起使用。 - BPL

现在好多了 ;) , +1 . 最后一件事，操作员应该小心使用多个装饰器，因为现在的方式强制用户在顶部指定 @transform_ast。 - BPL