BookReaderIA/datanode/windowed_iterator.py

   1 #!/usr/bin/python
   2
   3 # Copyright(c)2008-2010 Internet Archive. Software license AGPL version 3.
   4 #
   5 # This file is part of BookReader.
   6 #
   7 #     BookReader is free software: you can redistribute it and/or modify
   8 #     it under the terms of the GNU Affero General Public License as published by
   9 #     the Free Software Foundation, either version 3 of the License, or
  10 #     (at your option) any later version.
  11 #
  12 #     BookReader is distributed in the hope that it will be useful,
  13 #     but WITHOUT ANY WARRANTY; without even the implied warranty of
  14 #     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  15 #     GNU Affero General Public License for more details.
  16 #
  17 #     You should have received a copy of the GNU Affero General Public License
  18 #     along with BookReader.  If not, see <http://www.gnu.org/licenses/>.
  19 #
  20 #     The BookReader source is hosted at http://github.com/openlibrary/bookreader/
  21
  22 from collections import deque
  23 import itertools
  24
  25 class windowed_iterator:
  26     """ Wrap an iterator s.t. we can see [window] neighbors
  27     in either direction from the current item.
  28
  29     Items are stored in a deque of size 2*window + 1, where the latest
  30     item is always in the middle position.
  31
  32     The supplied clear_callback() is called for items more than
  33     [window] steps in the past.
  34
  35     """
  36
  37     # Todo? remove use of None as sentinel, to be able to represent
  38     # iterators returning None.
  39
  40     def __init__(self, iterator, window, clear_callback=None):
  41         self.iterator = iterator
  42         # initialize deque with sentinel values
  43         self.items = deque((None for i in range(window + 1)),
  44                            window * 2 + 1)
  45         self.window = window
  46         self.clear_callback = clear_callback
  47     def __iter__(self):
  48         return self
  49     def __repr__(self):
  50         return str(self.items) + ' window: ' + str(self.window)
  51     def clear(self):
  52         for item in self.items:
  53             if item and self.clear_callback is not None:
  54                 self.clear_callback(item)
  55         self.items.clear()
  56     def neighbor(self, delta):
  57         if abs(delta) > self.window:
  58             raise IndexError('Requested delta outside window')
  59         while self.window + delta + 1 > len(self.items):
  60             try:
  61                 self.items.append(self.iterator.next())
  62             except StopIteration:
  63                 return None
  64         return self.items[self.window + delta]
  65     def neighbors(self, window=None, modtwo=False):
  66         if window is None:
  67             window = self.window
  68         if window > self.window:
  69             raise IndexError('Requested delta outside window')
  70         for i in itertools.chain(range(-window, 0),
  71                                   range(1, window + 1)):
  72             if modtwo and i % 2 == 1:
  73                 continue
  74             n = self.neighbor(i)
  75             if n is not None:
  76                 yield n
  77     def next(self):
  78         nextitem = None
  79         if len(self.items) == self.window + 1:
  80             # elicit potential StopIteration before clearing/popping
  81             nextitem = self.iterator.next()
  82         if self.items[0] is not None and self.clear_callback is not None:
  83             self.clear_callback(self.items[0])
  84         self.items.popleft()
  85         if nextitem is not None:
  86             self.items.append(nextitem)
  87         return self.items[self.window]
  88
  89
  90 if __name__ == '__main__':
  91     def sample_gen():
  92         for i in range(0, 10):
  93             yield { 'num': i*i }
  94
  95     g = sample_gen()
  96     c = windowed_iterator(g, 3)
  97
  98     for i, item in enumerate(c):
  99         print 'item %s: %s' % (i, item)
 100         # print c
 101         if i in (1, 4, 6, 9):
 102             print 'neighbors of item %s: %s' % (i, [n for n in c.neighbors(2)])