src/allmydata/util/consumer.py

   1
   2 """This file defines a basic download-to-memory consumer, suitable for use in
   3 a filenode's read() method. See download_to_data() for an example of its use.
   4 """
   5
   6 from zope.interface import implements
   7 from twisted.internet.interfaces import IConsumer
   8
   9 class MemoryConsumer:
  10     implements(IConsumer)
  11
  12     def __init__(self, progress=None):
  13         self.chunks = []
  14         self.done = False
  15         self._progress = progress
  16
  17     def registerProducer(self, p, streaming):
  18         self.producer = p
  19         if streaming:
  20             # call resumeProducing once to start things off
  21             p.resumeProducing()
  22         else:
  23             while not self.done:
  24                 p.resumeProducing()
  25
  26     def write(self, data):
  27         self.chunks.append(data)
  28         if self._progress is not None:
  29             self._progress.set_progress(sum([len(c) for c in self.chunks]))
  30
  31     def unregisterProducer(self):
  32         self.done = True
  33
  34 def download_to_data(n, offset=0, size=None, progress=None):
  35     """
  36     :param on_progress: if set, a single-arg callable that receives total bytes downloaded
  37     """
  38     d = n.read(MemoryConsumer(progress=progress), offset, size)
  39     d.addCallback(lambda mc: "".join(mc.chunks))
  40     return d