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

Question

Let's say, I have a bunch of functions a, b, c, d and e and I want to find out if they directly use a loop:

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"

I want to write a function uses_loop so I can expect these assertions to pass:

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

(I expect uses_loop(c) to return False because c uses a list comprehension instead of a loop.)

I can't modify a, b, c, d and e. So I thought it might be possible to use ast for this and walk along the function's code which I get from inspect.getsource. But I'm open to any other proposals, this was only an idea how it could work.

This is as far as I've come with 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

You need to check if the function's Abstract Syntaxt Tree has any nodes that are an instance of ast.For or ast.While or ast.AsyncFor. You can use ast.walk() to visit every node of the 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)

See the documentation for ast, async for was added in 3.5.