How do I make a contextmanager with a loop inside?

Question

I want something like this:

from contextlib import contextmanager

@contextmanager
def loop(seq):
    for i in seq:
        try:
            do_setup(i)
            yield # with body executes here
            do_cleanup(i)
        except CustomError as e:
            print(e)

with loop([1,2,3]):
    do_something_else()
    do_whatever()

But contextmanager doesn't work because it expects the generator to yield exactly once.

The reason why I want this is because I basically want to make my own custom for loop. I have a modified IPython that is used to control test equipment. It's obviously a full Python REPL, but most of the time the user is just calling predefined functions (similar to Bash prompt), and the user is not expected to be a programmer or familiar with Python. There needs to be a way to loop over some arbitrary code with setup/cleanup and exception handling for each iteration, and it should be about as simple to type as the above with statement.

Simeon Visser · Accepted Answer · 2015-04-17T20:21:10.130

13

I think a generator works better here:

def loop(seq):
    for i in seq:
        try:
            print('before')
            yield i  # with body executes here
            print('after')
        except CustomError as e:
            print(e)

for i in loop([1,2,3]):
    print(i)
    print('code')

will give:

before
1
code
after
before
2
code
after
before
3
code
after

Python enters and exits a with block only once so you can't have logic int the enter / exit steps that would be done repeatedly.

edited Apr 17 '15 at 20:21

answered Apr 17 '15 at 20:17

Simeon Visser

118,920
18
185
180

An old question, but I can't see how the generator is equivalent to the context manager. ihmo, one of the main benefits of OP's code is that the try...except catches exceptions in the yielded-to code (eg. `do_whatever()`) .. this solution does not. Any Ideas on how to get that? Typical scenario would be some retry logic – ttyridal Feb 05 '19 at 08:26
@ttyridal: It depends where you raise the exception. If it happens in the generator, then the above is fine. If it happens in the `for i in loop...` outside, it won't get caught, and you need to catch it out there (maybe with a contextmanager). Also, contextlib.contextmanager creates a special generator that can only yield once, and I needed a thing that would yield for every element of the sequence. You need both a contextmanager and a generator to be a full solution (I posted this as another answer). – jpkotta Mar 11 '19 at 21:48

score 5 · Answer 2 · answered Mar 11 '19 at 21:49

5

A more complete answer, for if the exception might happen outside the generator:

from contextlib import contextmanager

class CustomError(RuntimeError):
    pass

@contextmanager
def handle_custom_error():
    try:
        yield
    except CustomError as e:
        print(f"handled: {e}")

def loop(seq):
    for i in seq:
        try:
            print('before')
            if i == 0:
                raise CustomError("inside generator")
            yield i # for body executes here
            print('after')
        except CustomError as e:
            print(f"handled: {e}")

@handle_custom_error()
def do_stuff(i):
    if i == 1:
        raise CustomError("inside do_stuff")
    print(f"i = {i}")

for i in loop(range(3)):
    do_stuff(i)

Output:

before
handled: inside generator
before
handled: inside do_stuff
after
before
i = 2
after

answered Mar 11 '19 at 21:49

jpkotta

9,237
3
29
34

That's a great solution. However it needs the `do_stuff` code to be in separate, decorated method. How about a solution that allows going with easier to use approach like this: `for i in loop(range(3)): if i == 1: raise CustomError("inside code")` Is there such a way to handle this in python without having to use both `for` with nested `with` in the user code? (I think ruby's code blocks would allow such a thing, hence my question) – Mirek Apr 04 '19 at 13:45
Basically I'm trying to enhance my test framework so the QA devs could easily use one simple structure to retry code (not just a method) that might raise AssertionError (or some other, custom exception tuple, thus the generator and context manager must share some scope to know which exception to handle). – Mirek Apr 04 '19 at 13:53
@Mirek I don't know if there's a way to do that. That's actually what I wanted to do in the first place but my answer is good enough for my purposes (all errors I want to catch are in predefined functions, and they can't really happen outside those functions). To do what you want, I might start by looking at how fuckit.py works. – jpkotta Apr 04 '19 at 21:37

How do I make a contextmanager with a loop inside?

2 Answers2

Linked