peekable iterator

peek.py

# using a tuple for self.buf was the fastest out of handful of
# different implementations I played around with.

class peekable(object):
    def __init__(self, iterable):
        """
        Extends iterators with a `peek` method that returns the next element
        that the iterator will yield. Raises StopIteration if there is no next
        element.

        >>> iterator = peekable(range(10))
        >>> while True:
        ...     try:
        ...         assert iterator.peek() == next(iterator)
        ...     except StopIteration:
        ...         break
        ...         
        """
        self.iter = iter(iterable)
        self.buf = ()

    def peek(self):
        if not self.buf:
            self.buf = (next(self.iter),)
        return self.buf[0]

    def __iter__(self):
        return self

    def __next__(self):
        if self.buf:
            val = self.buf[0]
            self.buf = ()
            return val
        return next(self.iter)