How to catch StopIteration from subgenerator

Question

I'd like to write a generator which can accept a limited number of inputs from yields and then gracefully handle further inputs. What's the best way of catching StopIteration?

I've tried wrapping by inner generator with an outer generator using a yield from expression inside a try-except block, but StopIteration gets raised anyway...

def limited_writer(max_writes):
    for i in range(max_writes):
        x = yield
        print(x) #

def graceful_writer(l):
    try:
        yield from l
    except StopIteration:
        # Ideally will have additional handling logic here
        raise Exception("Tried to write too much")

l_w = limited_writer(4)
g_w = graceful_writer(w)

g_w.send(None)

for i in range(5):
    g_w.send(i)

I'd like the above to raise Exception (but more generally provide a nice way of handling providing too much data), but in fact it still raises StopIteration. What's the best solution?

Answer 1

If you want graceful_writer to keep accepting data that is sent to it via its .send() method, it needs to keep on yielding indefinitely. The try / except block you currently have doesn't actually do anything, the yield from statement already absorbs the StopIteration from limited_writer . The one you are seeing at the top level comes from graceful_writer itself, when it reaches the end of its code.

To avoid that, try using an infinite loop, like this:

def graceful_writer(gen):
    yield from gen    # send values to wrapped generator for as long as it will take them
    while True:           
        yield         # then loop forever, discarding any additional values sent in