Search code examples
pythoncontextmanager

How do I make a contextmanager with a loop inside?


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.


Solution

  • 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.