pythonloopsabstract-syntax-treestatic-analysisinspect

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


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__)

Solution

  • Check if the function's abstract syntaxt tree (AST) has any ast.For or ast.While or ast.AsyncFor nodes. 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 for details. async for was added in Python 3.5.