Make compression_level configurable for .tar.gz (#888)

Author: gdh1995

pkg/private/tar/build_tar.py

@@ -42,7 +42,8 @@ class TarFile(object):
   class DebError(Exception):
     pass
 
-  def __init__(self, output, directory, compression, compressor, create_parents, allow_dups_from_deps, default_mtime):
+  def __init__(self, output, directory, compression, compressor, create_parents,
+               allow_dups_from_deps, default_mtime, compression_level):
     # Directory prefix on all output paths
     d = directory.strip('/')
     self.directory = (d + '/') if d else None
@@ -52,6 +53,7 @@ def __init__(self, output, directory, compression, compressor, create_parents, a
     self.default_mtime = default_mtime
     self.create_parents = create_parents
     self.allow_dups_from_deps = allow_dups_from_deps
+    self.compression_level = compression_level
 
   def __enter__(self):
     self.tarfile = tar_writer.TarFileWriter(
@@ -60,7 +62,8 @@ def __enter__(self):
         self.compressor,
         self.create_parents,
         self.allow_dups_from_deps,
-        default_mtime=self.default_mtime)
+        default_mtime=self.default_mtime,
+        compression_level=self.compression_level)
     return self
 
   def __exit__(self, t, v, traceback):
@@ -397,6 +400,9 @@ def main():
   parser.add_argument('--allow_dups_from_deps',
                       action='store_true',
                       help='')
+  parser.add_argument(
+      '--compression_level', default=-1,
+      help='Specify the numeric compress level in gzip mode; may be 0-9 or -1 (default to 6).')
   options = parser.parse_args()
 
   # Parse modes arguments
@@ -440,6 +446,10 @@ def main():
   if options.stamp_from:
     default_mtime = build_info.get_timestamp(options.stamp_from)
 
+  compression_level = -1
+  if options.compression_level:
+    compression_level = int(options.compression_level)
+
   # Add objects to the tar file
   with TarFile(
       options.output,
@@ -448,7 +458,8 @@ def main():
       compressor = options.compressor,
       default_mtime=default_mtime,
       create_parents=options.create_parents,
-      allow_dups_from_deps=options.allow_dups_from_deps) as output:
+      allow_dups_from_deps=options.allow_dups_from_deps,
+      compression_level = compression_level) as output:
 
     def file_attributes(filename):
       if filename.startswith('/'):

pkg/private/tar/tar.bzl

@@ -114,6 +114,8 @@ def _pkg_tar_impl(ctx):
                 "--owner_names",
                 "%s=%s" % (_quote(key), ctx.attr.ownernames[key]),
             )
+    if ctx.attr.compression_level:
+        args.add("--compression_level", ctx.attr.compression_level)
 
     # Now we begin processing the files.
     path_mapper = None
@@ -272,6 +274,10 @@ pkg_tar_impl = rule(
         ),
         "create_parents": attr.bool(default = True),
         "allow_duplicates_from_deps": attr.bool(default = False),
+        "compression_level": attr.int(
+            doc = """Specify the numeric compression level in gzip mode; may be 0-9 or -1 (default to 6).""",
+            default = -1,
+        ),
 
         # Common attributes
         "out": attr.output(mandatory = True),

pkg/private/tar/tar_writer.py

@@ -49,7 +49,8 @@ def __init__(self,
                create_parents=False,
                allow_dups_from_deps=True,
                default_mtime=None,
-               preserve_tar_mtimes=True):
+               preserve_tar_mtimes=True,
+               compression_level=-1):
     """TarFileWriter wraps tarfile.open().
 
     Args:
@@ -86,10 +87,11 @@ def __init__(self,
     else:
       mode = 'w:'
       if compression in ['tgz', 'gz']:
+        compression_level = min(compression_level, 9) if compression_level >= 0 else 6
         # The Tarfile class doesn't allow us to specify gzip's mtime attribute.
         # Instead, we manually reimplement gzopen from tarfile.py and set mtime.
         self.fileobj = gzip.GzipFile(
-            filename=name, mode='w', compresslevel=6, mtime=self.default_mtime)
+            filename=name, mode='w', compresslevel=compression_level, mtime=self.default_mtime)
     self.compressor_proc = None
     if self.compressor_cmd:
       mode = 'w|'

tests/tar/BUILD

@@ -477,6 +477,10 @@ py_test(
         ":test-tar-strip_prefix-substring.tar",
         ":test-tar-tree-artifact",
         ":test-tar-tree-artifact-noroot",
+        ":test-tar-compression_level--1",
+        ":test-tar-compression_level-3",
+        ":test-tar-compression_level-6",
+        ":test-tar-compression_level-9",
         ":test-tree-input-with-strip-prefix",
         ":test_tar_leading_dotslash",
         ":test_tar_package_dir_substitution.tar",
@@ -756,3 +760,12 @@ verify_archive_test(
     ],
     target = ":program_with_dir_runfiles_tar",
 )
+
+[pkg_tar(
+    name = "test-tar-compression_level-%s" % compression_level,
+    compression_level = compression_level,
+    deps = [
+        "//tests:testdata/tar_test.tar",
+    ],
+    extension = "tgz",
+) for compression_level in [-1, 3, 6, 9]]

tests/tar/pkg_tar_test.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """Testing for pkg_tar."""
 
+import os
 import tarfile
 import unittest
 
@@ -293,6 +294,17 @@ def test_externally_defined_duplicate_structure(self):
     ]
     self.assertTarFileContent('test-respect-externally-defined-duplicates.tar', content)
 
+  def test_compression_level(self):
+    sizes = [
+      ('test-tar-compression_level--1.tgz', 179),
+      ('test-tar-compression_level-3.tgz', 230),
+      ('test-tar-compression_level-6.tgz', 178),
+      ('test-tar-compression_level-9.tgz', 167),
+    ]
+    for file_name, expected_size in sizes:
+      file_path = runfiles.Create().Rlocation('rules_pkg/tests/tar/' + file_name)
+      file_size = os.stat(file_path).st_size
+      self.assertEqual(file_size, expected_size, 'size error for ' + file_name)
 
 if __name__ == '__main__':
   unittest.main()