Search code examples
pythonpython-3.xexceptiongeneratorstopiteration

Understanding StopIteration handling inside generators for non-trivial case


I am helping maintain some code that now includes automated Python 3.7 testing. This led me to some issues related to PEP 479 "Change StopIteration handling inside generators". My naive understanding was that you could use a try-except block to modify old code to be compatible with all python versions, e.g.

Old code:

def f1():
    it = iter([0])
    while True:
        yield next(it)

print(list(f1()))
# [0] (in Py 3.6)
# "RuntimeError: generator raised StopIteration" (in Py 3.7;
# or using from __future__ import generator_stop)

Becomes:

def f2():
    it = iter([0])
    while True:
        try:
            yield next(it)
        except StopIteration:
            return 

print(list(f2()))
# [0] (in all Python versions)

For this trivial example, it works, but I have found for some more complex code I am re-factoring it does not. Here is a minimal example with Py 3.6:

class A(list):
    it = iter([0])
    def __init__(self):
        while True:
            self.append(next(self.it))

class B(list):
    it = iter([0])
    def __init__(self):
        while True:
            try:
                self.append(next(self.it))
            except StopIteration:
                raise

class C(list):
    it = iter([0])
    def __init__(self):
        while True:
            try:
                self.append(next(self.it))
            except StopIteration:
                return  # or 'break'

def wrapper(MyClass):
    lst = MyClass()
    for item in lst:
        yield item

print(list(wrapper(A)))
# [] (wrong output)
print(list(wrapper(B)))
# [] (wrong output)
print(list(wrapper(C)))
# [0] (desired output)

I know that the A and B examples are exactly equivalent and that the C case is the correct way compatible with Python 3.7 (I also know that re-factoring to a for loop would make sense for many examples, including this contrived one).

But the question is why do the examples with A and B produce an empty list [], rather than [0]?


Solution

  • The first two cases have an uncaught StopIteration raised in the class's __init__. The list constructor handles this just fine in Python 3.6 (with possibly a warning, depending on the version). However, the exception propagates before wrapper gets a chance to iterate: the line that effectively fails is lst = MyClass(), and the loop for item in lst: never runs, causing the generator to be empty.

    When I run this code in Python 3.6.4, I get the following warning on both print lines (for A and B):

    DeprecationWarning: generator 'wrapper' raised StopIteration
    

    The conclusion here is twofold:

    1. Don't let the iterator run out on its own. It's your job to check when it stops. This is easy to do with a for loop, but has to be done manually with a while loop. Case A is a good illustration.
    2. Don't re-raise the internal exception. Return None instead. Case B is just not the way to go. A break or return would work correctly in the except block, as you did in C.

    Given that for loops are syntactic sugar for the try-except block in C, I would generally recommend their use, even with manual invocations of iter:

    class D(list):
        it = iter([0])
        def __init__(self):
            for item in it:
                self.append(item)
    

    This version is functionally equivalent to C, and does all the bookkeeping for you. There are very few cases that require an actual while loop (skipping calls to next being one that comes to mind, but even those cases can be rewritten with a nested loop).