p2pool/util/datachunker.py

   1 import collections
   2
   3 class StringBuffer(object):
   4     'Buffer manager with great worst-case behavior'
   5
   6     def __init__(self, data=''):
   7         self.buf = collections.deque([data])
   8         self.buf_len = len(data)
   9         self.pos = 0
  10
  11     def __len__(self):
  12         return self.buf_len - self.pos
  13
  14     def add(self, data):
  15         self.buf.append(data)
  16         self.buf_len += len(data)
  17
  18     def get(self, wants):
  19         if self.buf_len - self.pos < wants:
  20             raise IndexError('not enough data')
  21         data = []
  22         while wants:
  23             seg = self.buf[0][self.pos:self.pos+wants]
  24             self.pos += len(seg)
  25             while self.buf and self.pos >= len(self.buf[0]):
  26                 x = self.buf.popleft()
  27                 self.buf_len -= len(x)
  28                 self.pos -= len(x)
  29
  30             data.append(seg)
  31             wants -= len(seg)
  32         return ''.join(data)
  33
  34 def _DataChunker(receiver):
  35     wants = receiver.next()
  36     buf = StringBuffer()
  37
  38     while True:
  39         if len(buf) >= wants:
  40             wants = receiver.send(buf.get(wants))
  41         else:
  42             buf.add((yield))
  43 def DataChunker(receiver):
  44     '''
  45     Produces a function that accepts data that is input into a generator
  46     (receiver) in response to the receiver yielding the size of data to wait on
  47     '''
  48     x = _DataChunker(receiver)
  49     x.next()
  50     return x.send