Merge lp:~jderose/filestore/callback into lp:filestore

[Jason Gerard DeRose]

This changes batch_import_iter() to take an optional keyword-only `callback` argument. The callback takes 2 arguments:

def callback(count, size):
    pass

That's the completed count, and the completed bytes. The callback is called at the end of processing each file, plus whenever the next 128 MiB (16 leaves) have been completed in a large file and at least one more leaf remains.

For example, say you have a batch with one file and the file is exactly 33 leaves in size (264 MiB). The callback would be called 3 times, like this:

callback(0, 128 * MiB)
callback(0, 256 * MiB)
callback(1, 264 * MiB)

Updating the progress every 128 MiB is probably about right as far as reassuring the user, but we should play with it a bit if it seems too infrequent or unnecessarily frequent.

[David Jordan (dmj726)]

Looks good. I'm glad to see this UX issue fixed.

review: Approve

=== modified file 'filestore.py'
--- filestore.py	2011-11-28 21:27:39 +0000
+++ filestore.py	2011-12-08 09:31:33 +0000
@@ -740,7 +740,7 @@
         queue.put(e)
 
 
-def batch_import_iter(batch, *filestores):
+def batch_import_iter(batch, *filestores, callback=None):
     """
     Import the files in *batch* into one or more destination *filestores*.
 
@@ -769,6 +769,7 @@
     :param filestores: One or more `FileStore` instances in which to import the
         files in *batch*
     """
+    # Validate args:
     if not isinstance(batch, Batch):
         raise TypeError(TYPE_ERROR.format('batch', Batch, type(batch), batch))
     if not filestores:
@@ -779,6 +780,13 @@
         if not isinstance(fs, FileStore):
             name = 'filestores[{}]'.format(i)
             raise TypeError(TYPE_ERROR.format(name, FileStore, type(fs), fs))
+    if not (callback is None or callable(callback)):
+        raise TypeError('callback not callable: {!r}'.format(callback))
+
+    # Okay, time to work:
+    if callback is not None:
+        count = 0
+        size = 0
     q = SmartQueue(16)
     thread = _start_thread(batch_reader, batch, q)
     while True:
@@ -788,6 +796,9 @@
         if file.size == 0:
             if q.get() is not EndFile:
                 raise Exception('file.size is 0, but not followed by EndFile')
+            if callback is not None:
+                count += 1
+                callback(count, size)
             yield BatchItem(file, None)
             continue
         temps = tuple(fs.allocate_tmp(file.size) for fs in filestores)
@@ -796,6 +807,8 @@
             leaf = q.get()
             if leaf is EndFile:
                 break
+            if callback is not None and leaf.index and leaf.index % 16 == 0:
+                callback(count, size + h.file_size)
             h.hash_leaf(leaf)
             for tmp_fp in temps:
                 tmp_fp.write(leaf.data)
@@ -805,6 +818,10 @@
         for (i, fs) in enumerate(filestores):
             tmp_fp = temps[i]
             fs.move_to_canonical(tmp_fp, ch.id)
+        if callback is not None:
+            count += 1
+            size += file.size
+            callback(count, size)
         yield BatchItem(file, ch)
     thread.join()  # Make sure batch_reader() terminates
 
 
=== modified file 'test_filestore.py'
--- test_filestore.py	2011-11-23 04:25:45 +0000
+++ test_filestore.py	2011-12-08 09:31:33 +0000
@@ -136,6 +136,14 @@
         self.items.append(item)
 
 
+class Callback:
+    def __init__(self):
+        self.calls = []
+
+    def __call__(self, *args):
+        self.calls.append(args)
+
+
 class TestSmartQueue(TestCase):
     def test_get(self):
         q = filestore.SmartQueue()
@@ -1094,6 +1102,14 @@
             str(cm.exception),
             'batch_import_iter() takes at least one destination FileStore'
         )
+        
+        # Test with non-callable callback
+        with self.assertRaises(TypeError) as cm:
+            list(filestore.batch_import_iter(batch, fs1, fs2, callback='hello'))
+        self.assertEqual(
+            str(cm.exception),
+            "callback not callable: 'hello'"
+        )
 
         # Test with wrong filestore type
         with self.assertRaises(TypeError) as cm:
@@ -1142,6 +1158,62 @@
             '[Errno 13] Permission denied: {!r}'.format(src.name)
         )
 
+        # Test with callback and empty batch
+        cb = Callback()
+        batch = filestore.Batch(tuple(), 0, 0)
+        tmp1 = TempDir()
+        fs1 = filestore.FileStore(tmp1.dir)
+        tmp2 = TempDir()
+        fs2 = filestore.FileStore(tmp2.dir)
+        self.assertEqual(
+            list(filestore.batch_import_iter(batch, fs1, fs2, callback=cb)),
+            []
+        )
+        self.assertEqual(list(fs1), [])
+        self.assertEqual(list(fs2), [])
+        self.assertEqual(cb.calls, [])
+
+        # Test with callback and non-empty batch
+        tmp = TempDir()
+        src = filestore.File(tmp.join('src.mov'), FILE_SIZE, 0)
+        write_sample_file(src.name)
+        empty1 = filestore.File(tmp.touch('empty1.mov'), 0, 0)
+        empty2 = filestore.File(tmp.touch('empty2.mov'), 0, 0)
+
+        file_size = filestore.LEAF_SIZE * 33
+        big = filestore.File(tmp.join('big.mov'), file_size, 0)
+        data = b'a' * filestore.LEAF_SIZE
+        fp = open(big.name, 'wb')
+        for i in range(33):
+            fp.write(data)
+        fp.close()
+        fp = open(big.name, 'rb')
+        ch = filestore.hash_fp(fp)
+
+        batch = filestore.Batch((big, empty1, src, empty2), 0, 0)
+        self.assertEqual(
+            list(filestore.batch_import_iter(batch, fs1, fs2, callback=cb)),
+            [
+                (big, ch),
+                (empty1, None),
+                (src, CH),
+                (empty2, None),
+            ]
+        )
+        self.assertEqual(set(s.id for s in fs1), set([ID, ch.id]))
+        self.assertEqual(set(s.id for s in fs2), set([ID, ch.id]))
+        self.assertEqual(
+            cb.calls,
+            [
+                (0, filestore.LEAF_SIZE * 16),  # 128 MiB done
+                (0, filestore.LEAF_SIZE * 32),  # 256 MiB done
+                (1, file_size),
+                (2, file_size),
+                (3, file_size + 20971520),
+                (4, file_size + 20971520)
+            ]
+        )
+
     def test_hash_fp(self):
         tmp = TempDir()
         src = tmp.join('foo.mov')