Question

As part of some WSGI middleware I want to write a python class that wraps an iterator to implement a close method on the iterator.

This works fine when I try it with an old-style class, but throws a TypeError when I try it with a new-style class. What do I need to do to get this working with a new-style class?

Example:

class IteratorWrapper1:

    def __init__(self, otheriter):
        self._iterator = otheriter
        self.next = otheriter.next

    def __iter__(self):
        return self

    def close(self):
        if getattr(self._iterator, 'close', None) is not None:
            self._iterator.close()
        # other arbitrary resource cleanup code here

class IteratorWrapper2(object):

    def __init__(self, otheriter):
        self._iterator = otheriter
        self.next = otheriter.next

    def __iter__(self):
        return self

    def close(self):
        if getattr(self._iterator, 'close', None) is not None:
            self._iterator.close()
        # other arbitrary resource cleanup code here

if __name__ == "__main__":
    for i in IteratorWrapper1(iter([1, 2, 3])):
        print i

    for j in IteratorWrapper2(iter([1, 2, 3])):
        print j

Gives the following output:

1
2
3
Traceback (most recent call last):
  ...
TypeError: iter() returned non-iterator of type 'IteratorWrapper2'

Answer 1

What you’re trying to do makes sense, but there’s something evil going on inside Python here.

class foo(object):
    c = 0
    def __init__(self):
        self.next = self.next2

    def __iter__(self):
        return self

    def next(self):
        if self.c == 5: raise StopIteration
        self.c += 1
        return 1

    def next2(self):
        if self.c == 5: raise StopIteration
        self.c += 1
        return 2

it = iter(foo())
# Outputs: <bound method foo.next2 of <__main__.foo object at 0xb7d5030c>>
print it.next
# 2
print it.next()
# 1?!
for x in it:
    print x

foo() is an iterator which modifies its next method on the fly–perfectly legal anywhere else in Python. The iterator we create, it, has the method we expect: it.next is next2. When we use the iterator directly, by calling next(), we get 2. Yet, when we use it in a for loop, we get the original next, which we’ve clearly overwritten.

I’m not familiar with Python internals, but it seems like an object’s “next” method is being cached in tp_iternext (http://docs.python.org/c-api/typeobj.html#tp_iternext), and then it’s not updated when the class is changed.

This is definitely a Python bug. Maybe this is described in the generator PEPs, but it’s not in the core Python documentation, and it’s completely inconsistent with normal Python behavior.

You could work around this by keeping the original next function, and wrapping it explicitly:

class IteratorWrapper2(object):
    def __init__(self, otheriter):
        self.wrapped_iter_next = otheriter.next
    def __iter__(self):
        return self
    def next(self):
        return self.wrapped_iter_next()

for j in IteratorWrapper2(iter([1, 2, 3])):
    print j

… but that’s obviously less efficient, and you should not have to do that.