Python – Learn about StopIteration handling inside generators in non-trivial cases

Learn about StopIteration handling inside generators in non-trivial cases… here is a solution to the problem.

Learn about StopIteration handling inside generators in non-trivial cases

I’m helping maintain some code that now includes automated Python 3.7 tests. This brings me to some issues related to PEP 479 “Change StopIteration handling inside generators”. My naïve understanding is that you can use try-except block to modify old code to be compatible with all python versions, eg

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 simple example it works, but I found that for some more complex code I was refactoring it and it didn’t. This is a minimal example of 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 the same, and the C example is the right way to be compatible with Python 3.7 (I also know that refactoring into a for loop makes sense for many examples, including this human-designed example).

But the question is why the example with A and B generates an empty list [], instead of [0]?

Solution

The first two cases raise an uncaptured StopIteration in the __init__ of the class. The list constructor handles well in Python 3.6 (there may be warnings, depending on the version). However, exception propagation before wrapper has a chance to iterate: the line that effectively fails is lst = MyClass(), and the loop for item in lst: never runs, resulting in an empty generator.

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 is your job to check when it stops. It’s easy to do this using a for loop, but it must be done manually using a while loop. Case A is a good example.
  2. Do not rethrow the inner exception. Instead, it returns None. Case B is not a viable approach. break or return will work fine in except block, just like you do in C.

Given that for loops are syntactic sugar for try-except blocks in C, I would generally recommend using them, even if iter:: is called manually

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 work for you. There are very few cases where an actual while loop is required (skipping the call to next is the one that comes to mind, but even these cases can be overridden with nested loops).

Related Problems and Solutions