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

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) 到 return 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

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

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

Answer 2

你快到了！您所要做的就是找出如何从正文中获取数据 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

Answer 3

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

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

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

python

loops

static-analysis

abstract-syntax-tree

inspect