Best way to receive the 'return' value from a python generator
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?
Answer summary
Combining answers from @Chad S. and @KT, the simplest appears to turn my generator function into a class using the iterator protocol:
class ValueGenerator():
def __iter__(self):
yield 1
yield 2
# and so on
self.summary = {...}
vg = ValueGenerator()
for i in vg:
do_something(i)
values_summary = vg.summary
And @Ferdinand Beyer's answer is simplest if I can't refactor the value producer.
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
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
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)
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
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
>>>