Python通过函数装饰器转换ast

时间:2019-08-13 14:53:40

标签: python abstract-syntax-tree python-decorators

我有个主意,可以使用类似于下面的装饰器来转换所有已标记的给定功能,

@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个)

我在做错什么吗?

1 个答案:

答案 0 :(得分: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`

将解决问题。

其他参考:Exploring Python Code Objects