Welcome, guest | Sign In | My Account | Store | Cart

Windowing an iterable with itertools (Python recipe) by Daniel Cohn
ActiveState Code (http://code.activestate.com/recipes/577196/)

Oftentimes a programmer needs to peek into an iterator without advancing it, a task for which many good solutions already exist. But what if the intrepid coder needs a fast and pythonic way to 'window' the data? This recipe demonstrates how to wrap any iterable with a class that adds two methods, prev and peek.

Python, 24 lines
Download
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
from itertools import tee

class Iterator(object):
    """Intended to be used inside a while loop"""
    def __init__(self, iterable):
        self._a, self._b = tee(iter(iterable), 2)
        self._previous = None
        self._peeked   = self._b.next()
    
    def __iter__(self):
        return self
    
    def next(self):
        self._previous = self._a.next()
        self._current  = self._peeked
        try:
            self._peeked = self._b.next()
        except StopIteration:
            self._peeked = None
        return self._current
        
    def prev(self): return self._previous
    
    def peek(self): return self._peeked

For example,

I = Iterator([1, 2, 3, 4])
fcn = lambda itr: (itr.prev(), itr.next(), itr.peek())

L = []
while True:
    try:
        L.append(fcn(I))
    except StopIteration:
        break

print 'L: ', L

returns [(None, 1, 2), (1, 2, 3), (2, 3, 4), (3, 4, None)]. Even cooler, instead of passing a list into Iterator we could have passed a generator! So what if we want to peek into the iterable using a 'for' loop? Well, the recipe can be modified as follows:

class Window(object):
    """Intended to be used with a for loop"""
    def __init__(self, iterable):
        self._a, self._b = tee(iter(iterable), 2)
        self._previous = None
        self._peeked   = self._b.next()

    def __iter__(self):
        return self

    def next(self):
        _prev = self._previous
        self._previous = self._a.next()
        self._current  = self._peeked
        try:
            self._peeked = self._b.next()
        except StopIteration:
            self._peeked = None
        return _prev, self._current, self._peeked

    def prev(self): return self._previous

    def peek(self): return self._peeked

Then the following:

W = Window([1, 2, 3, 4])
L = []

for prev, current, peeked in W:
    L.append((prev, current, peeked))

print 'L: ', L

returns [(None, 1, 2), (1, 2, 3), (2, 3, 4), (3, 4, None)].

Tags: generator, iterable, iterator, loop, peeking, window