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（）不接受任何参数（给定1个）

我做错了什么吗？

Answer 1

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

使用execmode编译的代码始终被视为模块级代码，当然，它可以包含函数或类定义，或任何其他有效的 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

因此，在创建 new 时，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 代码对象