Best way to receive the 'return' value from a python generator

Question

Since Python 3.3, if a generator function returns a value, that becomes the value for the StopIteration exception that is raised. This can be collected a number of ways:

• The value of a yield from expression, which implies the enclosing function is also a generator.
• Wrapping a call to next() or .send() in a try/except block.

However, if I'm simply wanting to iterate over the generator in a for loop - the easiest way - there doesn't appear to be a way to collect the value of the StopIteration exception, and thus the return value. Im using a simple example where the generator yields values, and returns some kind of summary at the end (running totals, averages, timing statistics, etc).

for i in produce_values():
    do_something(i)

values_summary = ....??

One way is to handle the loop myself:

values_iter = produce_values()
try:
    while True:
        i = next(values_iter)
        do_something(i)
except StopIteration as e:
    values_summary = e.value

But this throws away the simplicity of the for loop. I can't use yield from since that requires the calling code to be, itself, a generator. Is there a simpler way than the roll-ones-own for loop shown above?

Artificiality

I don't think generator return values were ever intended to be used outside the context of a subgenerator returning a value to an enclosing generator, so it makes sense that this would be awkward. An explicit try-except is probably your best bet. (Also, you want e.value, not just e.) – Personal

@user2357112: Fixed the coding error: thank you. I understand the comment, but it seems such a useful construct that I'm surprised it's awkward like that. The answer might well be "For loops are for iterators. if you're using a generator for anything other than a simple iterator, then for loops are Not For You." – Artificiality

I think it's more that you're trying to give the generator too many responsibilities. Things like summaries or timing information would more reasonably be done by the caller. If you really want to do something like this, I'd probably create a wrapper around the generator and give the wrapper a summary method, to be called once iteration completes. – Personal

@Personal : That'd only work if there was sufficient information in the values to form those summaries. THis is why I added "timing" to the summary information in my contrived example :) However, if I'm going to that level of detail then a class-with-iterator-protocol is probably far more appropriate. – Artificiality

The normal way to save state is to write your own iterator class. I think the statement would be "For loops are for sequences. If you want to create the sequence dynamically but also want to maintain state after the for loop exits, write your own iterator class instead of using a simple generator." – Leighleigha

(You got the class conversion wrong; making __next__ a generator function just makes your iterator return an endless stream of generators.) – Personal

@Personal Thanks... copy-paste-editing problem. Fixed. – Artificiality

Answer 1

You can think of the value attribute of StopIteration (and arguably StopIteration itself) as implementation details, not designed to be used in "normal" code.

Have a look at PEP 380 that specifies the yield from feature of Python 3.3: It discusses that some alternatives of using StopIteration to carry the return value where considered.

Since you are not supposed to get the return value in an ordinary for loop, there is no syntax for it. The same way as you are not supposed to catch the StopIteration explicitly.

A nice solution for your situation would be a small utility class (might be useful enough for the standard library):

class Generator:
    def __init__(self, gen):
        self.gen = gen

    def __iter__(self):
        self.value = yield from self.gen
        return self.value

This wraps any generator and catches its return value to be inspected later:

>>> def test():
...     yield 1
...     return 2
...
>>> gen = Generator(test())
>>> for i in gen:
...    print(i)
...
1
>>> print(gen.value)
2

The line return self.value in __iter__() ensures that the return value will be propagated correctly when Generator instances are nested (e.g. gen = Generator(Generator(test()))).

Humblebee

a simple and elegant extension to the syntax of the for loop that would carry such StopIteration value was proposed by a good colleague of mine a while back, but it was turned down... – Azotemia

It would be great if one could specify for i in test() as value to obtain value as the return value of the generator. – Bedfast

Word to the wise about this construction: also have the __iter__ method of Generator return self.value, in case you have nested Generators. (That is, an instance of Generator for which self.gen is itself a Generator.) If you don't do that, the return value of the inner instance won't pass to the outer instance. I got bitten by this. – Blanchard

@Bedfast What happens if there's a break statement inside the loop? Also, the scoping would confuse me. How about making use of the existing for/else syntax, where the else runs if the loop completes normally? e.g. for i in test(): ... else value: ... – Bartholemy

I would like to see an example of that syntax, @Azotemia – Immense

@EricAuld That's a good point, I've edited the answer to incorporate your suggestion because it doesn't cause any harm and it's more correct. – Metz

Answer 2

A light-weight way to handle the return value (one that doesn't involve instantiating an auxiliary class) is to use dependency injection.

Namely, one can pass in the function to handle / act on the return value using the following wrapper / helper generator function:

def handle_return(generator, func):
    returned = yield from generator
    func(returned)

For example, the following--

def generate():
    yield 1
    yield 2
    return 3

def show_return(value):
    print('returned: {}'.format(value))

for x in handle_return(generate(), show_return):
    print(x)

results in--

1
2
returned: 3

Thebes

Answer 3

You could make a helper wrapper, that would catch the StopIteration and extract the value for you:

from functools import wraps

class ValueKeepingGenerator(object):
    def __init__(self, g):
        self.g = g
        self.value = None
    def __iter__(self):
        self.value = yield from self.g

def keep_value(f):
    @wraps(f)
    def g(*args, **kwargs):
        return ValueKeepingGenerator(f(*args, **kwargs))
    return g

@keep_value
def f():
    yield 1
    yield 2
    return "Hi"

v = f()
for x in v:
    print(x)

print(v.value)

Storfer

This is probably the best answer if I was not able to modify the "produce_values" generator myself. If that was in my control, I'd likely have the function be a class following the iterator protocol. – Artificiality

You came up with a similar solution as I. But if I may: Mine is way simpler :) – Humblebee

@ChrisCogdon: I don't think you're likely to find any such produce_values generators you didn't write. It's not a pattern people commonly use. – Personal

@FerdinandBeyer: True. In the spirit of competition I edited the answer now by stealing your idea now and decorating it a bit. – Storfer

@Personal : very likely true 😸 – Artificiality

I'd suggest not initializing self.value to None because trying to access the value before iteration was finished should probably be an error. – Juggle

I think it is a matter of taste and personal preferences in error-handling. For example, when value is initialized you can check it even if the iteration fails for some reason. However, I see how not initializing it might be a more "common" approach, so will edit the answer now. – Storfer

@Storfer Nice idea of using a decorator! Maybe use functools.wraps to make it even better. – Humblebee

@FerdinandBeyer: Indeed, I was just trying to recall what package was this wraps decorator in. Thanks. – Storfer

For python >= 3.8, walrus operator works here: for x in (v := f()): print(x) then print(v.value) – Amboina

Answer 4

The most obvious method I can think of for this would be a user defined type that would remember the summary for you..

>>> import random
>>> class ValueProducer:
...    def produce_values(self, n):
...        self._total = 0
...        for i in range(n):
...           r = random.randrange(n*100)
...           self._total += r
...           yield r
...        self.value_summary = self._total/n
...        return self.value_summary
... 
>>> v = ValueProducer()
>>> for i in v.produce_values(3):
...    print(i)
... 
25
55
179
>>> print(v.value_summary)
86.33333333333333
>>> 

Juggle

Good idea. Additionally, rather than a method-generator, I could make the object itself an iterator by implementing a __next__ method. – Artificiality

Answer 5

Another light weight way sometimes appropriate is to yield the running summary in every generator step in addition to your primary value in a tuple. The loop stays simple with an extra binding which is still available afterwards:

for i, summary in produce_values():
    do_something(i)

show_summary(summary)

This is especially useful if someone could use more than just the last summary value, e. g. updating a progress view.

Snooze

Answer 6

Now with full type annotation:

import typing
T = typing.TypeVar("T")
U = typing.TypeVar("U")
V = typing.TypeVar("V")


class ValuedGenerator(typing.Generic[T, U, V]):
    value: V | None = None

    def __init__(self, generator: typing.Generator[T, U, V]):
        self.generator = generator
        super().__init__()

    def __iter__(self) -> typing.Generator[T, U, None]:
        self.value = yield from self.generator

Loni