do_bloom(): remove unused "count" variable

[bup.git] / cmd / bloom-cmd.py

diff --git a/cmd/bloom-cmd.py b/cmd/bloom-cmd.py
index 5e3cd8689dbda1350d71d7cd6197492c8869bd76..d74cce2704567597fc1fce934032dfd025041e55 100755 (executable)
--- a/cmd/bloom-cmd.py
+++ b/cmd/bloom-cmd.py
@@ -6,19 +6,52 @@ from bup.helpers import *
 optspec = """
 bup bloom [options...]
 --
+ruin       ruin the specified bloom file (clearing the bitfield)
 f,force    ignore existing bloom file and regenerate it from scratch
 o,output=  output bloom filename (default: auto)
 d,dir=     input directory to look for idx files (default: auto)
 k,hashes=  number of hash functions to use (4 or 5) (default: auto)
+c,check=   check the given .idx file against the bloom filter
 """
 
+
+def ruin_bloom(bloomfilename):
+    rbloomfilename = git.repo_rel(bloomfilename)
+    if not os.path.exists(bloomfilename):
+        log("%s\n" % bloomfilename)
+        add_error("bloom: %s not found to ruin\n" % rbloomfilename)
+        return
+    b = bloom.ShaBloom(bloomfilename, readwrite=True, expected=1)
+    b.map[16:16+2**b.bits] = '\0' * 2**b.bits
+
+
+def check_bloom(path, bloomfilename, idx):
+    rbloomfilename = git.repo_rel(bloomfilename)
+    ridx = git.repo_rel(idx)
+    if not os.path.exists(bloomfilename):
+        log("bloom: %s: does not exist.\n" % rbloomfilename)
+        return
+    b = bloom.ShaBloom(bloomfilename)
+    if not b.valid():
+        add_error("bloom: %r is invalid.\n" % rbloomfilename)
+        return
+    base = os.path.basename(idx)
+    if base not in b.idxnames:
+        log("bloom: %s does not contain the idx.\n" % rbloomfilename)
+        return
+    if base == idx:
+        idx = os.path.join(path, idx)
+    log("bloom: bloom file: %s\n" % rbloomfilename)
+    log("bloom:   checking %s\n" % ridx)
+    for objsha in git.open_idx(idx):
+        if not b.exists(objsha):
+            add_error("bloom: ERROR: object %s missing" 
+                      % str(objsha).encode('hex'))
+
+
 _first = None
 def do_bloom(path, outfilename):
     global _first
-    if not outfilename:
-        assert(path)
-        outfilename = os.path.join(path, 'bup.bloom')
-
     b = None
     if os.path.exists(outfilename) and not opt.force:
         b = bloom.ShaBloom(outfilename)
@@ -48,14 +81,14 @@ def do_bloom(path, outfilename):
 
     if b:
         if len(b) != rest_count:
-            log("bloom: size %d != idx total %d, regenerating\n"
-                    % (len(b), rest_count))
+            debug1("bloom: size %d != idx total %d, regenerating\n"
+                   % (len(b), rest_count))
             b = None
         elif (b.bits < bloom.MAX_BLOOM_BITS and
               b.pfalse_positive(add_count) > bloom.MAX_PFALSE_POSITIVE):
-            log("bloom: regenerating: adding %d entries gives "
-                "%.2f%% false positives.\n"
-                    % (add_count, b.pfalse_positive(add_count)))
+            debug1("bloom: regenerating: adding %d entries gives "
+                   "%.2f%% false positives.\n"
+                   % (add_count, b.pfalse_positive(add_count)))
             b = None
         else:
             b = bloom.ShaBloom(outfilename, readwrite=True, expected=add_count)
@@ -68,7 +101,7 @@ def do_bloom(path, outfilename):
     msg = b is None and 'creating from' or 'adding'
     if not _first: _first = path
     dirprefix = (_first != path) and git.repo_rel(path)+': ' or ''
-    log('bloom: %s%s %d file%s (%d object%s).\n'
+    progress('bloom: %s%s %d file%s (%d object%s).\n'
         % (dirprefix, msg,
            len(add), len(add)!=1 and 's' or '',
            add_count, add_count!=1 and 's' or ''))
@@ -76,18 +109,20 @@ def do_bloom(path, outfilename):
     tfname = None
     if b is None:
         tfname = os.path.join(path, 'bup.tmp.bloom')
-        tf = open(tfname, 'w+')
-        b = bloom.ShaBloom.create(tfname, f=tf, expected=add_count, k=opt.k)
-    count = 0
+        b = bloom.create(tfname, expected=add_count, k=opt.k)
+
     icount = 0
     for name in add:
         ix = git.open_idx(name)
         qprogress('bloom: writing %.2f%% (%d/%d objects)\r' 
                   % (icount*100.0/add_count, icount, add_count))
         b.add_idx(ix)
-        count += 1
         icount += len(ix)
 
+    # Currently, there's an open file object for tfname inside b.
+    # Make sure it's closed before rename.
+    b.close()
+
     if tfname:
         os.rename(tfname, outfilename)
 
@@ -100,12 +135,24 @@ o = options.Options(optspec)
 if extra:
     o.fatal('no positional parameters expected')
 
-if opt.k and opt.k not in (4,5):
-    o.fatal('only k values of 4 and 5 are supported')
-
 git.check_repo_or_die()
 
+if not opt.check and opt.k and opt.k not in (4,5):
+    o.fatal('only k values of 4 and 5 are supported')
+
 paths = opt.dir and [opt.dir] or git.all_packdirs()
 for path in paths:
     debug1('bloom: scanning %s\n' % path)
-    do_bloom(path, opt.output)
+    outfilename = opt.output or os.path.join(path, 'bup.bloom')
+    if opt.check:
+        check_bloom(path, outfilename, opt.check)
+    elif opt.ruin:
+        ruin_bloom(outfilename)
+    else:
+        do_bloom(path, outfilename)
+
+if saved_errors:
+    log('WARNING: %d errors encountered during bloom.\n' % len(saved_errors))
+    sys.exit(1)
+elif opt.check:
+    log('All tests passed.\n')