Friendly usage of a Python iterable over a sequence of context managers

Question

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.

Answer 1

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.