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

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

您需要检查函数的抽象语法树是否有任何节点是 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)

参见 documentation for ast , async for 是 added in 3.5 .