Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

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.

like image 654
Chris Cogdon Avatar asked Dec 03 '15 18:12

Chris Cogdon


People also ask

How do you get data from a generator in Python?

To get values from the generator object, call the next() method on the generator object or loop through the generator object.

Can you return in a generator function Python?

Python provides a generator to create your own iterator function. A generator is a special type of function which does not return a single value, instead, it returns an iterator object with a sequence of values. In a generator function, a yield statement is used rather than a return statement.

What can we use instead of return in Python?

Summary: The yield keyword in python works like a return with the only difference is that instead of returning a value, it gives back a generator function to the caller.

Are Python generators efficient?

In Python, yield is a keyword that turns a function into a generator. Unlike a list, a generator does not store values. Instead, it knows the current value and how to get the next one. This makes a generator memory-efficient.


2 Answers

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 
like image 110
Ferdinand Beyer Avatar answered Sep 22 '22 22:09

Ferdinand Beyer


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) 
like image 27
KT. Avatar answered Sep 21 '22 22:09

KT.