From f8438d9419175feb0fbcfcf0d07cadfb7ac19306 Mon Sep 17 00:00:00 2001
From: Campbell Barton <ideasman42@gmail.com>
Date: Fri, 19 Dec 2014 11:22:59 +0100
Subject: [PATCH] bam cli: add compression level support

---
 bam/blend/blendfile_pack.py  | 14 +++++++++++++-
 bam/cli.py                   | 21 +++++++++++++++++++--
 tests/batch/bam_pack_test.py |  7 ++++---
 3 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/bam/blend/blendfile_pack.py b/bam/blend/blendfile_pack.py
index 9c6a4b2..f5116ce 100755
--- a/bam/blend/blendfile_pack.py
+++ b/bam/blend/blendfile_pack.py
@@ -104,6 +104,7 @@ def pack(
         deps_remap=None, paths_remap=None, paths_uuid=None,
         # load every libs dep, not just used deps.
         all_deps=False,
+        compress_level=-1,
         # yield reports
         report=None,
 
@@ -342,7 +343,15 @@ def pack(
 
     elif mode == 'ZIP':
         import zipfile
-        with zipfile.ZipFile(blendfile_dst.decode('utf-8'), 'w', zipfile.ZIP_DEFLATED) as zip_handle:
+
+        # not awesome!
+        import zlib
+        assert(compress_level in range(-1, 10))
+        _compress_level_orig = zlib.Z_DEFAULT_COMPRESSION
+        zlib.Z_DEFAULT_COMPRESSION = compress_level
+        _compress_mode = zipfile.ZIP_DEFLATED (compress_level == 0) if zipfile.ZIP_STORED else zipfile.ZIP_DEFLATED
+
+        with zipfile.ZipFile(blendfile_dst.decode('utf-8'), 'w', _compress_mode) as zip_handle:
             for fn in path_temp_files:
                 yield report("  %s: %r -> <archive>\n" % (colorize("copying", color='blue'), fn))
                 zip_handle.write(
@@ -368,6 +377,9 @@ def pack(
                     _dbg(b"RELATIVE_FILE: " + os.path.relpath(dst, base_dir_dst))
                     """
 
+        zlib.Z_DEFAULT_COMPRESSION = _compress_level_orig
+        del _compress_level_orig, _compress_mode
+
         yield report("  %s: %r\n" % (colorize("written", color='green'), blendfile_dst))
     else:
         raise Exception("%s not a known mode" % mode)
diff --git a/bam/cli.py b/bam/cli.py
index 5c85a99..28ec94b 100755
--- a/bam/cli.py
+++ b/bam/cli.py
@@ -826,6 +826,7 @@ class bam_commands:
             output,
             all_deps=False,
             use_quiet=False,
+            compress_level=-1,
             ):
         # Local packing (don't use any project/session stuff)
         from .blend import blendfile_pack
@@ -844,6 +845,7 @@ class bam_commands:
                 output.encode('utf-8'),
                 'ZIP',
                 all_deps=all_deps,
+                compress_level=compress_level,
                 report=report,
                 ):
             pass
@@ -926,7 +928,10 @@ def init_argparse_common(
         use_json=False,
         use_all_deps=False,
         use_quiet=False,
+        use_compress_level=False,
         ):
+    import argparse
+
     if use_json:
         subparse.add_argument(
                 "-j", "--json", dest="json", action='store_true',
@@ -942,6 +947,17 @@ def init_argparse_common(
                 "-q", "--quiet", dest="use_quiet", action='store_true',
                 help="Suppress status output",
                 )
+    if use_compress_level:
+        class ChoiceToZlibLevel(argparse.Action):
+            def __call__(self, parser, namespace, value, option_string=None):
+                setattr(namespace, self.dest, {"default": -1, "fast": 1, "best": 9, "store": 0}[value[0]])
+
+        subparse.add_argument(
+                "-c", "--compress", dest="compress_level", nargs=1, default=-1, metavar='LEVEL',
+                action=ChoiceToZlibLevel,
+                choices=('default', 'fast', 'best', 'store'),
+                help="Compression level for resulting archive",
+                )
 
 
 def create_argparse_init(subparsers):
@@ -1126,7 +1142,7 @@ def create_argparse_pack(subparsers):
             help="Output file or a directory when multiple inputs are passed",
             )
 
-    init_argparse_common(subparse, use_all_deps=True, use_quiet=True)
+    init_argparse_common(subparse, use_all_deps=True, use_quiet=True, use_compress_level=True)
 
     subparse.set_defaults(
             func=lambda args:
@@ -1134,7 +1150,8 @@ def create_argparse_pack(subparsers):
                     args.paths,
                     args.output,
                     all_deps=args.all_deps,
-                    use_quiet=args.use_quiet),
+                    use_quiet=args.use_quiet,
+                    compress_level=args.compress_level),
                     )
 
 
diff --git a/tests/batch/bam_pack_test.py b/tests/batch/bam_pack_test.py
index e4ad03d..ac93c0e 100755
--- a/tests/batch/bam_pack_test.py
+++ b/tests/batch/bam_pack_test.py
@@ -68,7 +68,7 @@ del bpy
 # Ensure module path
 import os
 import sys
-path = os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "client", "cli"))
+path = os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", ".."))
 if path not in sys.path:
     sys.path.append(path)
 
@@ -170,12 +170,13 @@ def pack_blend_test(blendfile_src, log, blender_bin):
     argv = (
         "pack", blendfile_src,
         "--output", TEMP_ZIP,
+        "--compress=store",
         "--quiet",
         )
 
-    import bam
+    import bam.cli
     log.info("bam " + " ".join(argv))
-    bam.main(argv)
+    bam.cli.main(argv)
 
     # extract zip
     os.makedirs(TEMP_EXTRACT, exist_ok=True)