Ryan C. Thompson
Ryan C. Thompson

Reputation: 41990

How can I run the initialization code for a generator function immediately, rather than at the first call?

I have a generator function that goes something like this:

def mygenerator():
    next_value = compute_first_value() # Costly operation
    while next_value != terminating_value:
        yield next_value
        next_value = compute_next_value()

I would like the initialization step (before the while loop) to run as soon as the function is called, rather than only when the generator is first used. What is a good way to do this?

I want to do this because the generator will be running in a separate thread (or process, or whatever multiprocessing uses) and I won't be using the return for a short while, and the initialization is somewhat costly, so I would like it to do the initialization while I'm getting ready to use the values.

Upvotes: 18

Views: 7134

Answers (5)

Andrey Sboev
Andrey Sboev

Reputation: 7672

class mygenerator(object):
    def __init__(self):
        self.next_value = self.compute_first_value()
    def __iter__(self):
        return self
    def next(self):
        if self.next_value == self.terminating_value:
            raise StopIteration()
        return self.next_value

Upvotes: 16

Tadhg McDonald-Jensen
Tadhg McDonald-Jensen

Reputation: 21453

For my use case I used a modified version of @ncoghlan answer but wrapped in a factory function to decorate the generating function:

import collections, functools, itertools

def primed_generator(generating_function):
    @functools.wraps(generating_function)
    def get_first_right_away_wrapper(*args,**kw):
        "call the generator function, prime before returning"
        gen = generating_function(*args,**kw)
        assert isinstance(gen,collections.Iterator)
        first_value = next(gen)
        return itertools.chain((first_value,),gen)
    return get_first_right_away_wrapper

Then just decorate the function:

@primed_generator
def mygenerator():
    next_value = compute_first_value() # Costly operation
    while next_value != terminating_value:
        yield next_value
        next_value = compute_next_value()

and the first value will be calculated immediately, and the result is transparent.

Upvotes: 0

Jason A
Jason A

Reputation: 583

I needed something similar. This is what I landed on. Push the generator function into an inner and return it's call.

def mygenerator():
    next_value = compute_first_value()

    def generator():
        while next_value != terminating_value:
            yield next_value
            next_value = compute_next(next_value)

    return generator()

Upvotes: 18

ncoghlan
ncoghlan

Reputation: 41486

You can create a "preprimed" iterator fairly easily by using itertools.chain:

from itertools import chain

def primed(iterable):
    """Preprimes an iterator so the first value is calculated immediately
       but not returned until the first iteration
    """
    itr = iter(iterable)
    try:
        first = next(itr)  # itr.next() in Python 2
    except StopIteration:
        return itr
    return chain([first], itr)

>>> def g():
...     for i in range(5):
...         print("Next called")
...         yield i
...
>>> x = primed(g())
Next called
>>> for i in x: print(i)
...
0
Next called
1
Next called
2
Next called
3
Next called
4

Upvotes: 9

Jessica
Jessica

Reputation: 6957

I suppose you can yield None after that first statement is completed, then in your calling code:

gen = mygenerator()
next(gen) # toss the None
do_something(gen)

Upvotes: 5

Related Questions