Search code examples
pythoniterationgeneratoryield

"yield" in Python


I have a function called x that produces a generator like this:

a = 5
def x():
    global a
    if a == 3:
        raise Exception("Stop")
    a = a - 1
    yield a

Then in the python shell I call that function like this:

>>> print x().next()
>>> 4
>>> print x().next()
>>> 3
>>> print x().next()
>>> <python-input-112-f3c02bba26c8> in x()
          2     global a
          3     if a == 3:
    ----> 4         raise Exception
          5     a = a - 1
          6     yield a

    Exception:

However, when I call that function and assign it to a variable, it behaves differently:

>>> a = 5
>>> b = x()
>>> print b.next()
>>> 4
>>> print b.next()
>>> ----> 1 b.next()
    StopIteration:

How is that even possible? Shouldn't it print out 3 and raise StopIteration in the next iteration?

PS: I know that when I first call the function, the body does not run, just produces a generator. The point I didn't understand is that what changes if I call and assign it to a variable?


Solution

  • In your first example, you were creating a new generator each time:

    x().next()
    

    This starts the generator from the top, so the first statement. When a == 3, the exception is raised, otherwise the generator just yields and pauses.

    When you assigned your generator later on, the global a started at 5, the code then continued from where it left of until it ends or comes across another yield statement, then ended. When a generator function ends, it raises StopIteration.

    Let's break this down into steps:

    1. a = 5.
    2. You create new generator and call .next() on it. The following code is executed:

      global a
      if a == 3:  # False
          raise Exception("Stop")
      a = a - 1   # a is now 4
      yield a
      

      The generator is paused on the last line, and 4 is yielded.

    3. You create a new generator and call .next() on it. a is 4 at the start:

      global a
      if a == 3:  # False
          raise Exception("Stop")
      a = a - 1   # a is now 3
      yield a
      

      The generator is paused on the last line, and 3 is yielded.

    4. You create a new generator and call .next() on it. a is 3 at the start:

      global a
      if a == 3:  # True
          raise Exception("Stop")
      

      An exception is raised.

    5. You set a = 5 again.

    6. You create a new generator, store a reference in b and call .next() on it. The following code is executed:

      global a
      if a == 3:  # False
          raise Exception("Stop")
      a = a - 1   # a is now 4
      yield a
      

      The generator is paused on the last line, and 4 is yielded.

    7. You call .next() again on the same, existing generator referenced by b. The code resumes at the paused point.

      The function has no more code at that point, and returns. StopIteration is raised.

    If you were to use a loop instead, you'd see the difference better:

    >>> def looping(stop):
    ...    for i in looping(stop):
    ...        yield i
    ...
    >>> looping(3).next()
    0
    >>> looping(3).next()
    0
    

    Note how each time I create a new generator, the loop starts from the beginning. Store a reference however, and you'll notice it continue instead:

    >>> stored = looping(3)
    >>> stored.next()
    0
    >>> stored.next()
    1
    >>> stored.next()
    2
    >>> stored.next()
    Traceback (most recent call last):
      File "<stdin>", line 1, in <module>
    StopIteration
    

    During the loop, each time the yield expression is executed, the code is paused; calling .next() continues the function where it left of the previous time.

    The StopIteration exception is entirely normal; it is how generators communicate that they are done. A for loop looks for this exception to end the loop:

    >>> for i in looping(3):
    ...     print i
    ... 
    0
    1
    2