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

Posted 2023-02-16

【中文标题】如何找出函数（的源代码）是不是包含循环？

【英文标题】：How to find out if (the source code of) a function contains a loop?如何找出函数（的源代码）是否包含循环？

【发布时间】：2019-06-03 05:01:31

【问题描述】：

假设我有一堆函数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 并沿着我从@​​987654338@ 获得的函数代码进行操作。 但我对任何其他建议持开放态度，这只是一个想法。

据我所知ast：

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

【参考方案1】：

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

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

【讨论】：

对于 OP 似乎关心的列表理解，这将返回 True。

虽然列表推导在技术上包含循环，但 Python 通过缩进识别循环。所以也许寻找缩进就足够了？

【参考方案2】：

你快到了！您所要做的就是找出如何从 body 对象中获取数据。在所有某些节点类型之后，它们都是属性。我刚刚使用getattr(node, 'body', []) 来获取孩子，如果其中任何一个属于_ast.For 或_ast.While，则返回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

【参考方案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 是added in 3.5。