Search code examples
pythoniterator

Iterator that supports pushback


I'm looking for a why to convert a regular iterator into one that supports pushing items back into it. E.g.

item = next(my_iterator)
if went_too_far(item):
    my_iterator.pushback(item)
    break;

This is similar, but not identical to, an iterator that supports peek; with the latter, the above would look more like this:

if went_too_far(my_iterator.peek()):
    break
else:
    item = next(my_iterator)

Solution

  • class PushbackWrapper(object):
    
        def __init__(self, iterator):
            self.__dict__['_iterator'] = iterator
            self.__dict__['_pushed'] = []
    
        def next(self):
            if len(self._pushed):
                return self._pushed.pop()
            else:
                return self._iterator.next()
    
        def pushback(self, item):
            self._pushed.append(item)
    
        def __getattr__(self, attr):
            return getattr(self._iterator, attr)
    
        def __setattr__(self, attr, value):
            return setattr(self._iterator, attr, value)
    

    To use it:

    pushback_enabled_iterator = PushbackWrapper(original_iterator)
    
    item = next(pushback_enabled_iterator)
    if went_too_far(item):
        pushback_enabled_iterator.pushback(item)
        break;