Search code examples
pythonfor-loopabstract-syntax-tree

Is it possible to identify if we are in the middle of a loop?


Is it possible to identify if we are in the middle of a loop?

I want to make the class automatically identify if the variable came from an iterator or not.

I checked inspect module and didn't succeed. Maybe there is some scope identifier or some brilliant magic methods manipulation that can flag when a code runs inside a loop.

Here's an example to make myself clear. How to set inside_loop Student attribute?

class Student:
    def __init__(self, id):
        self.id = id
        self.inside_loop = ???
    def print_grade(self):
        if self.inside_loop:
            print("you are inside a loop")
        else:
            print("checking grade...")     

>>> S = Student(1)
>>> S.print_grade()
>>> checking grade...

>>> student_ids = [1, 2, 3]
>>> for student in student_ids:
>>>     S = Student(student) 
>>>     S.print_grade()
>>>       
>>> you are inside a loop
>>> you are inside a loop
>>> you are inside a loop

Thanks in advance.


Solution

  • I have advanced a little bit in this question. I used ast and inspect standard modules of python. Since we had some worrisome comments, I would appreciate if you point out possible bugs in the code below. Until now, it solves the original problem.

    With inspect.currentframe().f_back I get the frame where the class was instantiate. Then the frame is used to build the abstract tree. For each tree node, it is set the parent attribute. Finally, I select the node that has the same line number than the original frame class instantiation assignment and check if this node parent is a ast.For, ast.While, or a ast.AsyncFor.

    It does not work in a pure python shell. See discussion in get ast from python object in interpreter.

    import ast
    import inspect
    
    class Student:
    
        def __init__(self, id):
            self.id = id
            self.previous_frame = inspect.currentframe().f_back
            self.inside_loop = self._set_inside_loop()
        
        def _set_inside_loop(self):
        
            loops = ast.For, ast.While, ast.AsyncFor
            nodes = ast.parse(inspect.getsource(self.previous_frame))
        
            for node in ast.walk(nodes):
                try:
                    for child in ast.iter_child_nodes(node):
                        child.parent = node
                
                    instantiate_class = node.lineno == self.previous_frame.f_lineno
                    loop_parent = isinstance(node.parent, loops)
                
                    if instantiate_class & loop_parent:
                        return True
                except Exception:
                    continue
            else:
                return False
        
        def print_grade(self):
            if self.inside_loop:
                print("you are inside a loop")
            else:
                print("checking grade...")  
    

    Now we have:

    d = Student(4)
    d.print_grade()
    
    >>> checking grade...
    
    students = [1,2,3]
    for i in students:
        s = Student(i)
        s.print_grade()
    
    >>> you are inside a loop
    >>> you are inside a loop
    >>> you are inside a loop
    
    i = 1
    while i<4:
        s = Student(i)
        s.print_grade()
        i +=1
    
    >>> you are inside a loop
    >>> you are inside a loop
    >>> you are inside a loop