How do I make a contextmanager with a loop inside?

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.


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