I'm writing a small library that tries to provide a persistent queue for dispatching jobs. My persistence code provides a way to iterate over pending job descriptions; I would also like to guarantee that dispatched jobs eventually get marked as completed or failed.
To do so, I first implemented it so that my user can do:
for c in some_iterator_object:
with c as x:
...
I dislike this solution for several reasons. First of all, I want to grab a job description from my queue as a single operation (and fail if the queue is empty), so the acquisition is done by the __next__
method of the iterator, and the release in the __exit__
of the context manager.
To ensure that the context manager is called, my __next__
returns a wrapper class that cannot be substituted directly for the value, so it will throw a clear error if the user forgets to call the context manager.
Is there a way to collapse these two statements into a single one ? Ideally, i would like to let the user do
for x in some_iterator_object:
...
all while being able to intercept exceptions raised by the contents of the for block.
EDIT: I found out by experimenting that if i let an unfinished generator get garbage collected, the yield statement will raise an internal exception, so I can write something crude like
try:
...
success = False
yield val
success = True
...
finally:
if success:
...
But if I understand correctly, this depends on the garbage collector to run, and it seems to be an internal mechanism that I shouldn't really touch.
If you want your context managers to be entered automatically as they are being returned by the iterator, you can write your own iterator class like this:
class ContextManagersIterator:
def __init__(self, it):
self._it = iter(it)
self._last = None
def __iter__(self):
return self
def __next__(self):
self.__exit__(None, None, None)
item = next(self._it)
item.__enter__()
self._last = item
return item
def __enter__(self):
return self
def __exit__(self, exc_type, exc_value, exc_traceback):
last = self._last
if last is not None:
self._last = None
return last.__exit__(exc_type, exc_value, exc_traceback)
Example usage:
from contextlib import contextmanager
@contextmanager
def my_context_manager(name):
print('enter', name)
try:
yield
finally:
print('exit', name)
sequence = [
my_context_manager('x'),
my_context_manager('y'),
my_context_manager('z'),
]
with ContextManagersIterator(sequence) as it:
for item in it:
print(' work')
# Output:
# enter x
# work
# exit x
# enter y
# work
# exit y
# enter z
# work
# exit z
This ContextManagersIterator
class takes care of calling __enter__
on its values just before they are returned. __exit__
is called right before another value is returned (if everything went well) or when an exception has been raised in the loop.