How to find out if (the source code of) a function contains a loop?
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.
You were almost there! All you had to do was to find out how to get the data from the body objects. They are all attributes after all of some Node type. I just used getattr(node, 'body', [])
to get the children and if any of them are of _ast.For
or _ast.While
return a True.
Note: I was just tinkering around the code. Not sure if this is documented somewhere and can be relied upon. I guess may be you can look it up? :)
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