Search code examples
pythonexceptiontry-catchyieldcontextmanager

How can a Python context manager try to execute code?


I'm trying to write a small context manager that'll try to execute some code repeatedly until the code works or until a specified number of tries has been made. I have attempted to write this but am encountering a difficulty with having the context manager handle problems when yielding:

Exception RuntimeError: 'generator ignored GeneratorExit'

How should I code this?

import contextlib
import random

def main():

    with nolube():
        print(1 / random.randint(0, 1))

@contextlib.contextmanager
def nolube(
    tries = None # None: try indefinitely
    ):
    """
    Create a context for trying something repeatedly.
    """
    tries_done = 0
    rekt = True
    if tries is None:
        while rekt is True:
            try:
                yield
                rekt = False
            except:
                tries_done += 1
                pass
    else:
        while rekt is True and tries_done <= tries:
            try:
                yield
                rekt = False
            except:
                tries_done += 1
                pass

if __name__ == "__main__":
    main()

Solution

  • @contextlib.contextmanager has a very clear contract; it'll only be resumed once. It can't be used to re-run code.

    In fact, you can't use a context manager to control repetitions at all. You need a loop here, not a context manager. A context manager doesn't control the block, it is only informed when entering and exiting.

    Use the tenacity package* instead; it provides a decorator. The decorator wraps a function in a while True loop that'll re-run the function for you.

    You'd apply it to your case by moving the print() statement into a function, decorated with @retry, then calling that function:

    import random
    from tenacity import retry
    
    @retry
    def foo():
        print(1 / random.randint(0, 1))
    
    def main():
        foo()
    

    * This answer originally recommended the retrying package but this was forked into a new package with updated API when that project fell dormant.