如何通过查看AST来告知生成器的普通Python函数?

时间:2017-05-02 10:37:31

标签: python python-3.x abstract-syntax-tree

我需要检测Python 3 AST中的ast.FunctionDef是正常的函数定义还是生成器定义。

我是否需要遍历身体并寻找ast.Yield - s或者是否有更简单的方法?

2 个答案:

答案 0 :(得分:3)

有一种偷偷摸摸的方法是用compile编译AST实例。代码对象附有几个标记,其中一个是'GENERATOR',您可以使用它们来区分这些标记。当然这取决于某些编译标志,因此它在CPython版本或实现中并不是真正可移植的

例如,使用非生成器函数:

func = """
def spam_func():
    print("spam")
"""
# Create the AST instance for it
m = ast.parse(func)
# get the function code
# co_consts[0] is used because `m` is 
# compiled as a module and we want the 
# function object
fc = compile(m, '', 'exec').co_consts[0]
# get a string of the flags and 
# check for membership
from dis import pretty_flags
'GENERATOR' in pretty_flags(fc.co_flags)  # False

同样,对于spam_gen生成器,您可以获得:

gen = """
def spam_gen():
    yield "spammy"
"""    
m = ast.parse(gen)
gc = compile(m, '', 'exec').co_consts[0]
'GENERATOR' in pretty_flags(gc.co_flags)  # True

这可能比你需要的更偷偷摸摸,穿越AST是另一个可行的选择,可能更容易理解和便携。

如果你有一个函数对象而不是AST,你总是可以使用func.__code__.co_flags执行相同的检查:

def spam_gen():
    yield "spammy"

from dis import pretty_flags
print(pretty_flags(spam_gen.__code__.co_flags))
# 'OPTIMIZED, NEWLOCALS, GENERATOR, NOFREE'

答案 1 :(得分:3)

遍历AST会比看起来更难 - 使用编译器可能是要走的路。这是一个为什么寻找Yield节点并不像听起来那么简单的例子。

$client->setAccessType("offline");
>>> s1 = 'def f():\n  yield'
>>> any(isinstance(node, ast.Yield) for node in ast.walk(ast.parse(s1)))
True
>>> dis.pretty_flags(compile(s1, '', 'exec').co_consts[0].co_flags)
'OPTIMIZED, NEWLOCALS, GENERATOR, NOFREE'

AST方法可能需要使用NodeVisitor来排除函数和lambda体。