Title: peeking into an iterator
Submitter:  Unknown or non-registered user
Last Updated: 2004/09/16
Version no: 1.2
Category:

Description:

Provides two different ways to look at the items in an iterator without removing them.

The peek function returns two values: the items retrieved and the reconstructed iterator.
The peekable class wraps an iterable as an iterator that provides a peek() method in addition to the usual next().

Source: Text Source

import itertools
import collections

def peek(iterable, n=None):
    """Returns the first item(s) of the iterable and an iterable equivalent
    to the original (i.e. without the first item(s) removed).

    Example usage:
    >>> peek(range(5))
    (0, <itertools.chain object at ...>)
    >>> first, iterable = peek(range(5))
    >>> first
    0
    >>> list(iterable)
    [0, 1, 2, 3, 4]
    >>> first, iterable = peek(range(5), 3)
    >>> first
    [0, 1, 2]
    >>> list(iterable)
    [0, 1, 2, 3, 4]
    >>> first, iterable = peek(range(5), 1)
    >>> first
    [0]
    >>> list(iterable)
    [0, 1, 2, 3, 4]
    """
    iterable = iter(iterable)
    if n is None:
        peek_result = iterable.next()
        iter_result = [peek_result]
    else:
        peek_result = iter_result = [iterable.next() for i in range(n)]
    return peek_result, itertools.chain(iter_result, iterable)


class peekable:
    """An iterator that supports a peek operation.

    Example Usage:
    >>> p = peekable(range(4))
    >>> p.peek()
    0
    >>> p.next(1)
    [0]
    >>> p.peek(3)
    [1, 2, 3]
    >>> p.next(2)
    [1, 2]
    >>> p.peek(2)
    Traceback (most recent call last):
      ...
    StopIteration
    >>> p.peek(1)
    [3]
    >>> p.next(2)
    Traceback (most recent call last):
      ...
    StopIteration
    >>> p.next()
    3
    """
    def __init__(self, iterable):
        self._iterable = iter(iterable)
        self._cache = collections.deque()
        
    def __iter__(self):
        return self

    def _fillcache(self, n):
        while len(self._cache) < n:
            self._cache.append(self._iterable.next())
    
    def next(self, n=None):
        self._fillcache(n is None and 1 or n)
        if n is None:
            result = self._cache.popleft()
        else:
            result = [self._cache.popleft() for i in range(n)]
        return result

    def peek(self, n=None):
        self._fillcache(n is None and 1 or n)
        if n is None:
            result = self._cache[0]
        else:
            result = [self._cache[i] for i in range(n)]
        return result

Discussion:

In my recipe for starzip (http://aspn.activestate.com/ASPN/Cookbook/Python/Recipe/302325), I needed to be able to peek at the first element of an iterator without actually removing the element from the iterator. The solution I used for that class is the peek function above. For starzip, the peek function was a reasonable solution -- extract the desired items and reconstruct the iterator using itertools.chain.

Recently, I was working on a simple parsing problem and again I wanted to be able to peek at the elements of an iterator. But this time, I needed to do multiple peeks, at many different points in the iterator. The peek function would have been unwieldy here -- not only would i have had to constantly reassign to my iterator variable, but I would have wrapped the iterator repeatedly in itertools.chain calls, which could mean substantial overhead because of the repeated function calls.

The solution provided by the peekable class avoids the repeated function calls by keeping a cache of items read from the iterator (implemented here as a deque). A side benefit of this approach is that if too many items are requested from the iterator, raising the StopIteration does not throw away the last values of the iterator. For example, while a StopIteration will be raised if next(5) is called when only 3 items are left in the iterator, a call later to next(3) will still be able to return these 3 items.

(The peekable solution is something akin to Danny Yoo's ExtendedIter in http://aspn.activestate.com/ASPN/Cookbook/Python/Recipe/197530, but allows for a larger lookahead and substantially simplifies the code by using a deque.)

[14 Sept 2004 -- Removed some comments to condense code.]
[16 Sept 2004 -- Changed default argument as per Peter Otten's suggestion.]