Flesh out "tahoe magic-folder status" command

[tahoe-lafs/tahoe-lafs.git] / src / allmydata / util / consumer.py

diff --git a/src/allmydata/util/consumer.py b/src/allmydata/util/consumer.py
index 4128c2006029e17a37146c95509285916b67b39e..a48fb59da9d5af3be9b198764b7ce51e807e8a1f 100644 (file)
--- 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