Send values to Python coroutine without handling StopIteration

Question

Given a Python coroutine:

def coroutine():
     score = 0
     for _ in range(3):
          score = yield score + 1

I'd like to use it in a simple loop like this:

cs = coroutine()
for c in cs:
     print(c)
     cs.send(c + 1)

... which I would expect to print

1
3
5

But actually, I get an exception on the line yield score + 1:

 TypeError: unsupported operand type(s) for +: 'NoneType' and 'int'

I can get it to work if I call next manually:

c = next(cs)
while True:
    print(c)
    try:
        c = cs.send(c + 1)
    except StopIteration:
        break

But I don't like that I need to use try/except, given that generators are usually so elegant.

So, is there any way to use a finite coroutine like this without explicitly handling StopIteration? I'm happy to change both the generator and the way I'm iterating over it.

Second Attempt

Martijn points out that both the for loop and my call to send advance the iterator. Fair enough. Why, then, can't I get around that with two yield statements in the coroutine's loop?

def coroutine():
    score = 0
    for _ in range(3):
        yield score
        score = yield score + 1

cs = coroutine()
for c in cs:
    print(c)
    cs.send(c + 1)

If I try that, I get the same error but on the send line.

0
None
Traceback (most recent call last):
  File "../coroutine_test.py", line 10, in <module>
    cs.send(c + 1)
TypeError: unsupported operand type(s) for +: 'NoneType' and 'int'

Answer 1

I'll take a stab at your second attempt. First, let coroutine be defined as:

def coroutine():
    score = 0
    for _ in range(3):
        yield
        score = yield score + 1

This function will output your 1, 3, 5 as in the original question.

Now, let's convert the for loop into a while loop.

# for loop
for c in cs:
    print(c)
    cs.send(c + 1)

# while loop
while True:
    try:
        c = cs.send(None)
        print(c)
        cs.send(c + 1)
    except StopIteration:
        break

Now, we can get this while loop working using the following if we precede it with a next(cs). In total:

cs = coroutine()
next(cs)
while True:
    try:
        c = cs.send(None)
        print(c)
        cs.send(c + 1)
    except StopIteration:
        break
# Output: 1, 3, 5

When we try to convert this back into a for loop, we have the relatively simple code:

cs = coroutine()
next(cs)
for c in cs:
    print(c)
    cs.send(c + 1)

And this outputs the 1, 3, 5 as you wanted. The issue is that in the last iteration of the for loop, cs is already exhausted, but send is called again. So, how do we get another yield out of the generator? Let's add one to the end...

def coroutine():
    score = 0
    for _ in range(3):
        yield
        score = yield score + 1
    yield

cs = coroutine()
next(cs)
for c in cs:
    print(c)
    cs.send(c + 1)
# Output: 1, 3, 5

This final example iterates as intended without a StopIteration exception.

Now, if we take a step back, this can all be better written as:

def coroutine():
    score = 0
    for _ in range(3):
        score = yield score + 1
        yield # the only difference from your first attempt

cs = coroutine()
for c in cs:
    print(c)
    cs.send(c + 1)
# Output: 1, 3, 5

Notice how the yield moved, and the next(cs) was removed.

Answer 2

Yes, with coroutines you generally have to use a next() call first to 'prime' the generator; it'll cause the generator function to execute code until the first yield. Your issue is mostly that you are using a for loop, however, which uses next() as well, but doesn't send anything.

You could add an extra yield to the coroutine to catch that first priming step, and add the @consumer decorator from PEP 342 (adjusted for Python 2 and 3):

def consumer(func):
    def wrapper(*args,**kw):
        gen = func(*args, **kw)
        next(gen)
        return gen
    wrapper.__name__ = func.__name__
    wrapper.__dict__ = func.__dict__
    wrapper.__doc__  = func.__doc__
    return wrapper

@consumer
def coroutine():
    score = 0
    yield
    for _ in range(3):
        score = yield score + 1

You'd still have to use a while loop, as a for loop can't send:

c = 0
while True:
    try:
        c = cs.send(c + 1)
    except StopIteration:
        break
    print(c)

Now, if you want this to work with a for loop, you have to understand when the next() call from the for loop comes in when you are in the loop. When the .send() resumes the generator, the yield expression returns the sent value, and the generator continues on from there. So the generator function only stops again the next time a yield appears.

So looking at a loop like this:

for _ in range(3):
    score = yield score + 1

the first time you use send the above code has already executed yield score + 1 and that'll now return the sent value, assigning it to score. The for loop continues on and takes the next value in the range(3), starts another iteration, then executes yield score + 1 again and pauses at that point. It is that next iteration value that is then produced.

Now, if you want to combine sending with plain next() iteration, you can add extra yield expressions, but those then need to be positioned such that your code is paused in the right locations; at a plain yield value when you are going to call next() (because it'll return None) and at a target = yield when you are using generator.send() (because it'll return the sent value):

@consumer
def coroutine():
    score = 0
    yield  # 1
    for _ in range(3):
        score = yield score + 1  # 2
        yield  # 3

When you use the above '@consumer' decorated generator with a for loop, the following happens:

• the @consumer decorator 'primes' the generator by moving to point 1.
• the for loop calls next() on the generator, and it advances to point 2, producing the score + 1 value.
• a generator.send() call returns paused generator at point 2, assigning the sent value to score, and advances the generator to point 3. This returns None as the generator.send() result!
• the for loop calls next() again, advancing to point 2, giving the loop the next score + 1 value.
• and so on.

So the above works directly with your loop:

>>> @consumer
... def coroutine():
...     score = 0
...     yield  # 1
...     for _ in range(3):
...         score = yield score + 1  # 2
...         yield  # 3
...
>>> cs = coroutine()
>>> for c in cs:
...     print(c)
...     cs.send(c + 1)
...
1
3
5

Note that the @consumer decorator and the first yield can now go again; the for loop can do that advancing to point 2 all by itself:

def coroutine():
    score = 0
    for _ in range(3):
        score = yield score + 1  # 2, for advances to here
        yield  # 3, send advances to here

and this still continues to work with your loop:

>>> def coroutine():
...     score = 0
...     for _ in range(3):
...         score = yield score + 1  # 2, for advances to here
...         yield  # 3, send advances to here
...
>>> cs = coroutine()
>>> for c in cs:
...     print(c)
...     cs.send(c + 1)
...
1
3
5