adding zlib and zlib-blocks tests

diff --git a/testing/run_benchmark.py b/testing/run_benchmark.py
index 48daf3d..9884a46 100755 (executable)
--- a/testing/run_benchmark.py
+++ b/testing/run_benchmark.py
@@ -34,7 +34,9 @@ def main():
     benchmark test.
     '''
     parser = argparse.ArgumentParser(description='Profiling test options. ')
-    parser.add_argument('-t', '--test', choices=['delta-tarfile', 'tarfile'],
+    parser.add_argument('-t', '--test',
+                        choices=['delta-tarfile', 'tarfile', 'zlib',
+                                 'zlib-blocks'],
                         help='Select option for testing. ')
     parser.add_argument('-l', '--compression-level', type=int, default=9,
                         choices=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
@@ -45,6 +47,10 @@ def main():
                         help='Enable profile')
     parser.add_argument('-m', '--tarmode',  default="w:gz",
                         help='Enable profile')
+    parser.add_argument('-st', '--strategy', default='default',
+                        choices=['default', 'filtered', 'huffman_only'],
+                        help='Select the deflate strategy in zlib related '
+                             'tests.')
     parser.add_argument(
         '-s', '--sort',
         help='Sort output profile',
@@ -64,6 +70,10 @@ def main():
 
     if pargs.test in ['delta-tarfile', 'tarfile']:
         test_tarfile(pargs.test, compression_level, path, pargs.tarmode)
+    elif pargs.test == 'zlib':
+        test_zlib(compression_level, path, pargs.strategy)
+    elif pargs.test == 'zlib-blocks':
+        test_zlib_blocks(compression_level, path, pargs.strategy)
     else:
         parser.print_help()
 
@@ -90,6 +100,94 @@ def test_tarfile(res, compression_level=9, path='source_dir', tarmode="w:gz"):
                            compresslevel=compression_level)
         tar.add(path)
 
+def test_zlib(compression_level=9, path='source_dir', strategy='default',
+              BUFSIZE = 16 * 1024):
+    '''
+    Test that compresses a directory with zlib
+    '''
+    if strategy == 'default':
+        strategy = zlib.Z_DEFAULT_STRATEGY
+    elif strategy ==  'filtered':
+        strategy = zlib.Z_FILTERED
+    else:
+        strategy = zlib.Z_HUFFMAN_ONLY
+
+    c = zlib.compressobj(compression_level,
+                            zlib.DEFLATED,
+                            -zlib.MAX_WBITS,
+                            zlib.DEF_MEM_LEVEL,
+                                strategy)
+
+    def add(path):
+        l = 0
+        l_orig = 0
+        if os.path.isdir(path):
+            for f in os.listdir(path):
+                l2, l_orig2 = add(os.path.join(path, f))
+                l += l2
+                l_orig += l_orig2
+        elif os.path.isfile(path):
+            with open(path, 'rb') as fo:
+                while True:
+                    buf = fo.read(BUFSIZE)
+                    l_orig += len(buf)
+                    if len(buf) == 0:
+                        return l, l_orig
+                    buf = c.compress(buf)
+                    if buf is not None:
+                        l += len(buf)
+        return l, l_orig
+
+    l, l_orig = add(path)
+    l += len(c.flush())
+    print("total size: %d MiB compressed, %d MiB uncompressed" % (
+        (l / (1024 ** 2)),
+        (l_orig / (1024 ** 2))))
+
+
+def test_zlib_blocks(compression_level=9, path='source_dir', strategy='default',
+                     BUFSIZE = 16 * 1024):
+    '''
+    Test that compresses a directory with zlib
+    '''
+    if strategy == 'default':
+        strategy = zlib.Z_DEFAULT_STRATEGY
+    elif strategy ==  'filtered':
+        strategy = zlib.Z_FILTERED
+    else:
+        strategy = zlib.Z_HUFFMAN_ONLY
+
+    def add(path):
+        l = 0
+        l_orig = 0
+        if os.path.isdir(path):
+            for f in os.listdir(path):
+                l2, l_orig2 = add(os.path.join(path, f))
+                l += l2
+                l_orig += l_orig2
+        elif os.path.isfile(path):
+            c = zlib.compressobj(compression_level,
+                                zlib.DEFLATED,
+                                -zlib.MAX_WBITS,
+                                zlib.DEF_MEM_LEVEL,
+                                strategy)
+            with open(path, 'rb') as fo:
+                while True:
+                    buf = fo.read(BUFSIZE)
+                    l_orig += len(buf)
+                    if len(buf) == 0:
+                        l += len(c.flush())
+                        return l, l_orig
+                    buf = c.compress(buf)
+                    if buf is not None:
+                        l += len(buf)
+        return l, l_orig
+
+    l, l_orig = add(path)
+    print("total size: %d MiB compressed, %d MiB uncompressed" % (
+        (l / (1024 ** 2)),
+        (l_orig / (1024 ** 2))))
+
 def print_profile(pr, sort):
     s = io.StringIO()
     if not sort: