Search code examples
pythongeneratoryield

How can I tell whether a generator was just-started?


I'd like a function, is_just_started, which behaves like the following:

>>> def gen(): yield 0; yield 1
>>> a = gen()
>>> is_just_started(a) 
True
>>> next(a)
0
>>> is_just_started(a) 
False
>>> next(a)
1
>>> is_just_started(a) 
False
>>> next(a)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration
>>> is_just_started(a)
False

How can I implement this function?

I looked at the .gi_running attribute but it appears to be used for something else.

If I know the first value that needs to be sent into the generator, I can do something like this:

def safe_send(gen, a):
    try:
        return gen.send(a)
    except TypeError as e:
        if "just-started" in e.args[0]:
            gen.send(None)
            return gen.send(a)
        else:
            raise

However, this seems abhorrent.


Solution

  • This only works in Python 3.2+:

    >>> def gen(): yield 0; yield 1
    ... 
    >>> a = gen()
    >>> import inspect
    >>> inspect.getgeneratorstate(a)
    'GEN_CREATED'
    >>> next(a)
    0
    >>> inspect.getgeneratorstate(a)
    'GEN_SUSPENDED'
    >>> next(a)
    1
    >>> inspect.getgeneratorstate(a)
    'GEN_SUSPENDED'
    >>> next(a)
    Traceback (most recent call last):
      File "<stdin>", line 1, in <module>
    StopIteration
    >>> inspect.getgeneratorstate(a)
    'GEN_CLOSED'
    

    So, the requested function is:

    import inspect
    
    def is_just_started(gen):
        return inspect.getgeneratorstate(gen) == inspect.GEN_CREATED:
    

    Out of curiosity, I looked into CPython to figure out how it was determining this... Apparently it looks at generator.gi_frame.f_lasti which is the "index of last attempted instruction in bytecode". If it's -1 then it hasn't started yet.

    Here's a py2 version:

    def is_just_started(gen):
        return gen.gi_frame is not None and gen.gi_frame.f_lasti == -1