如何找出函数的(源代码)是否包含循环？

Question

比方说,我有一堆的功能a,b,c,d和e我想看看他们是否直接使用一个循环:

def a():
    for i in range(3):
        print(i**2)

def b():
    i = 0
    while i < 3:
        print(i**2)
        i += 1

def c():
    print("\n".join([str(i**2) for i in range(3)]))

def d():
    print("\n".join(["0", "1", "4"]))

def e():
    "for"

我想写一个函数,uses_loop所以我可以期望这些断言传递:

assert uses_loop(a) == True
assert uses_loop(b) == True
assert uses_loop(c) == False
assert uses_loop(d) == False
assert uses_loop(e) == False

(我希望uses_loop(c)返回,False因为c使用列表解析而不是循环.)

我不能修改a,b,c,d和e.所以我认为可以使用ast这个并沿着我得到的函数代码inspect.getsource.但我对任何其他提案持开放态度,这只是一个想法,它是如何工作的.

就我而言ast:

def uses_loop(function):
    import ast
    import inspect
    nodes = ast.walk(ast.parse(inspect.getsource(function)))
    for node in nodes:
        print(node.__dict__)

Answer 1

您需要检查函数的摘要Syntaxt树有是的一个实例的任何节点ast.For或ast.While或ast.AsyncFor.您可以使用ast.walk()访问AST的每个节点

import ast
import inspect

def uses_loop(function):
    loop_statements = ast.For, ast.While, ast.AsyncFor

    nodes = ast.walk(ast.parse(inspect.getsource(function)))
    return any(isinstance(node, loop_statements) for node in nodes)

请参阅有关文档ast,async for是在3.5加入.