Combine and speed up idx->midx and bupindex merge

[bup.git] / lib / bup / helpers.py

diff --git a/lib/bup/helpers.py b/lib/bup/helpers.py
index 36b95602e9c8f2088ec9fd55d9693f47619b1d4f..8a4da97056caf558a3d1aa39c0898f59329b1b4f 100644 (file)
--- a/lib/bup/helpers.py
+++ b/lib/bup/helpers.py
@@ -1,6 +1,7 @@
 """Helper functions and classes for bup."""
 
 import sys, os, pwd, subprocess, errno, socket, select, mmap, stat, re
+import heapq, operator
 from bup import _version
 
 # This function should really be in helpers, not in bup.options.  But we
@@ -87,6 +88,36 @@ def next(it):
         return None
 
 
+def merge_iter(iters, pfreq, pfunc, pfinal, key=None):
+    if key:
+        samekey = lambda e, pe: getattr(e, key) == getattr(pe, key, None)
+    else:
+        samekey = operator.eq
+    count = 0
+    total = sum(len(it) for it in iters)
+    iters = (iter(it) for it in iters)
+    heap = ((next(it),it) for it in iters)
+    heap = [(e,it) for e,it in heap if e]
+
+    heapq.heapify(heap)
+    pe = None
+    while heap:
+        if not count % pfreq:
+            pfunc(count, total)
+        e, it = heap[0]
+        if not samekey(e, pe):
+            pe = e
+            yield e
+        count += 1
+        try:
+            e = it.next() # Don't use next() function, it's too expensive
+        except StopIteration:
+            heapq.heappop(heap) # remove current
+        else:
+            heapq.heapreplace(heap, (e, it)) # shift current to new location
+    pfinal(count, total)
+
+
 def unlink(f):
     """Delete a file at path 'f' if it currently exists.