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