X-Git-Url: https://git.rkrishnan.org/?a=blobdiff_plain;f=src%2Fallmydata%2Futil%2Fconsumer.py;h=a48fb59da9d5af3be9b198764b7ce51e807e8a1f;hb=144d31b4c355d87a1a4b7cf5ed5b3283fdca792e;hp=4128c2006029e17a37146c95509285916b67b39e;hpb=f660aa78ab02b90b3b8edf7521ed2a2515e0e7f2;p=tahoe-lafs%2Ftahoe-lafs.git

diff --git a/src/allmydata/util/consumer.py b/src/allmydata/util/consumer.py
index 4128c200..a48fb59d 100644
--- a/src/allmydata/util/consumer.py
+++ b/src/allmydata/util/consumer.py
@@ -8,9 +8,12 @@ from twisted.internet.interfaces import IConsumer
 
 class MemoryConsumer:
     implements(IConsumer)
-    def __init__(self):
+
+    def __init__(self, progress=None):
         self.chunks = []
         self.done = False
+        self._progress = progress
+
     def registerProducer(self, p, streaming):
         self.producer = p
         if streaming:
@@ -19,12 +22,19 @@ class MemoryConsumer:
         else:
             while not self.done:
                 p.resumeProducing()
+
     def write(self, data):
         self.chunks.append(data)
+        if self._progress is not None:
+            self._progress.set_progress(sum([len(c) for c in self.chunks]))
+
     def unregisterProducer(self):
         self.done = True
 
-def download_to_data(n, offset=0, size=None):
-    d = n.read(MemoryConsumer(), offset, size)
+def download_to_data(n, offset=0, size=None, progress=None):
+    """
+    :param on_progress: if set, a single-arg callable that receives total bytes downloaded
+    """
+    d = n.read(MemoryConsumer(progress=progress), offset, size)
     d.addCallback(lambda mc: "".join(mc.chunks))
     return d