Reusing Python generator objects

Generators are one of those little things that I love about Python. Never heard of generators? Well, this is all much better explained here, but in a nutshell:

    def squares( size ):
        for x in xrange( size ):
            yield x * x
        return

The above code returns a generator object. This object is iterable, and each item is returned by the yield statement. The return statement ends the iteration.

    >>> list( squares( 10 ) )
    [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]
    >>> for value in squares( 5 ):
    ...     print value
    ...
    0
    1
    4
    9
    16

Now, there’s a little problem I found… I wanted an object than, when iterated, returned a certain sequence of objects. No problem so far. But here’s the catch: once it finished returning the list, I wanted to be able to reuse the generator.

    >>> gen = squares( 5 )
    >>> list( gen )        # Works the first time...
    [0, 1, 4, 9, 16]
    >>> list( gen )        # But not the second!
    []

So I came up with this solution:

    class Regenerator(object):
        """
        Calls a generator and iterates it. When it's finished iterating, the
        generator is called again. This allows you to iterate a generator more
        than once (well, sort of, since the generator is actually recreated).
        """
        def __iter__(self):
            return self
        def __init__(self, g_function, *v_args, **d_args):
            self.__g_function = g_function
            self.__v_args     = v_args
            self.__d_args     = d_args
            self.__g_object   = None
        def next(self):
            if self.__g_object is None:
                self.__g_object = self.__g_function( *self.__v_args, **self.__d_args )
            try:
                return self.__g_object.next()
            except StopIteration:
                self.__g_object = None
                raise

So, in the example we presented, this is what we’d do to be able to iterate the list of numbers more than once:

    def _squares( size ):
        for x in xrange( size ):
            yield x * x
        return
    def squares( size ):
        return Regenerator( _squares, size )

And alas, it works! 🙂

    >>> gen = squares( 10 )
    >>> list( gen )
    [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]
    >>> list( gen )
    [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]
    >>> list( gen )
    [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

Download

regenerator.py