How to check if a python function uses a while loop?

user_10973705 0
Viewed 560
def foo():
    while <condition>:
        do something

def bar():
    for i in range(5):
        do something

Suppose I have two functions defined above in a file name test.py. Is there a way in python to write a function with the following behavior?

import test

def uses_while(fn: Callable) -> bool:
    (what goes here?)

>>> uses_while(test.foo)
True
>>> uses_while(test.bar)
False

I essentially need to programatically check if a function makes use of a while loop, without inspecting the code by hand. I've thought of using pdb.getsourcelines(), but that doesn't work if there are comments or strings with the word 'while' inside. Any ideas?

automated-tests bytecode python-3.x unit-testing python
2 Answers
import ast
import inspect
from typing import Callable

def uses_while(fn: Callable) -> bool:
    nodes = ast.walk(ast.parse(inspect.getsource(fn)))
    return any(isinstance(node, ast.While) for node in nodes)

On Python 3.9+ you'll have to change it to from collections.abc import Callable.

user_109737050

I wrote a simple function that can check if the function given as parameter contains a while loop:

import inspect

def test_while(func):
  flag = False

  body = inspect.getsourcelines(func)
  string = ''.join(body[0]).replace(' ', '')
  splited = string.split('\n')
  
  for chain in splited:
    if len(chain) > 0 and chain[0] is not '#':
      if chain.startswith('while'):
        flag = True

  return flag
user_101683210
Related