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

How to find out if (the source code of) a function contains a loop?

比方说,我有一堆函数 abcde,我想知道是否他们直接使用循环:

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) 到 return False 因为 c 使用列表理解而不是循环。)

我无法修改 abcde。所以我认为可以为此使用 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__)

如果您只是想检查函数体是否包含关键字 'for' 或 'while',您可以执行以下操作:

def uses_loop(func_name):
    import inspect
    lines = inspect.getsource(func_name)
    return 'for' in lines or 'while' in lines

你快到了!您所要做的就是找出如何从正文中获取数据 objects。毕竟它们都是某种节点类型的属性。我刚刚使用 getattr(node, 'body', []) 得到了 children,如果它们中的任何一个属于 _ast.For_ast.While return,则为 True。

注意:我只是在修改代码。不确定这是否记录在某处并且可以依赖。我想也许你可以查一下? :)

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 uses_loop(function):
    import ast
    import _ast
    import inspect
    nodes = ast.walk(ast.parse(inspect.getsource(function)))
    return any(isinstance(node, (_ast.For, _ast.While)) for node in nodes)


print(uses_loop(a))    # True
print(uses_loop(b))    # True
print(uses_loop(c))    # False
print(uses_loop(d))    # False

您需要检查函数的抽象语法树是否有任何节点是 ast.Forast.Whileast.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 was added in 3.5