Support custom compression program

pkg/archive.py

@@ -121,6 +121,7 @@ class Error(Exception):
   def __init__(self,
                name,
                compression='',
+               compressor='',
                root_directory='.',
                default_mtime=None,
                preserve_tar_mtimes=True):
@@ -129,6 +130,7 @@ def __init__(self,
     Args:
       name: the tar file name.
       compression: compression type: bzip2, bz2, gz, tgz, xz, lzma.
+      compressor: custom command to do the compression.
       root_directory: virtual root to prepend to elements in the archive.
       default_mtime: default mtime to use for elements in the archive.
           May be an integer or the value 'portable' to use the date
@@ -148,6 +150,7 @@ def __init__(self,
       else:
         self.use_xz_tool = True
     self.name = name
+    self.compressor = compressor
     self.root_directory = root_directory.rstrip('/').rstrip('\\')
     self.root_directory = self.root_directory.replace('\\', '/')
     self.preserve_mtime = preserve_tar_mtimes
@@ -422,7 +425,15 @@ def close(self):
     # Close the gzip file object if necessary.
     if self.fileobj:
       self.fileobj.close()
-    if self.use_xz_tool:
+    if self.compressor:
+      print(self.compressor)
+      if subprocess.call(
+          '< {0} {1} > {0}.d && mv {0}.d {0}'.format(self.name,
+                                                     self.compressor),
+          shell=True,
+          stdout=subprocess.PIPE):
+        raise self.Error("Compression with '{}' failed".format(self.compressor))
+    elif self.use_xz_tool:
       # Support xz compression through xz... until we can use Py3
       if subprocess.call('which xz', shell=True, stdout=subprocess.PIPE):
         raise self.Error('Cannot handle .xz and .lzma compression: '

pkg/build_tar.py

@@ -29,18 +29,20 @@ class TarFile(object):
   class DebError(Exception):
     pass
 
-  def __init__(self, output, directory, compression, root_directory,
+  def __init__(self, output, directory, compression, compressor, root_directory,
                default_mtime):
     self.directory = directory
     self.output = output
     self.compression = compression
+    self.compressor = compressor
     self.root_directory = root_directory
     self.default_mtime = default_mtime
 
   def __enter__(self):
     self.tarfile = archive.TarFileWriter(
         self.output,
         self.compression,
+        self.compressor,
         self.root_directory,
         default_mtime=self.default_mtime)
     return self
@@ -233,8 +235,14 @@ def main():
   parser.add_argument(
       '--directory',
       help='Directory in which to store the file inside the layer')
-  parser.add_argument('--compression',
-                      help='Compression (`gz` or `bz2`), default is none.')
+
+  compression = parser.add_mutually_exclusive_group()
+  compression.add_argument('--compression',
+                           help='Compression (`gz` or `bz2`), default is none.')
+  compression.add_argument('--compressor',
+                           help='Compressor program and arguments, '
+                                'e.g. `pigz -p 4`')
+
   parser.add_argument(
       '--modes', action='append',
       help='Specific mode to apply to specific file (from the file argument),'
@@ -298,7 +306,8 @@ def main():
   # Add objects to the tar file
   with TarFile(
       options.output, helpers.GetFlagValue(options.directory),
-      options.compression, options.root_directory, options.mtime) as output:
+      options.compression, options.compressor, options.root_directory,
+      options.mtime) as output:
 
     def file_attributes(filename):
       if filename.startswith('/'):

pkg/docs/reference.md

@@ -68,8 +68,9 @@ There are currently no other well-known attributes.
 ## pkg_tar
 
 ```python
-pkg_tar(name, extension, strip_prefix, package_dir, srcs,
-        mode, modes, deps, symlinks, package_file_name, package_variables)
+pkg_tar(name, extension, strip_prefix, package_dir, srcs, compressor,
+        compressor_args, mode, modes, deps, symlinks, package_file_name,
+        package_variables)
 ```
 
 Creates a tar file from a list of inputs.
@@ -145,6 +146,20 @@ Creates a tar file from a list of inputs.
         </p>
       </td>
     </tr>
+    <tr>
+      <td><code>compressor</code></td>
+      <td>
+        <code>Label, optional</code>
+        <p>Executable to be used as compression program.</p>
+      </td>
+    </tr>
+    <tr>
+      <td><code>compressor_args</code></td>
+      <td>
+        <code>String, optional</code>
+        <p>Arguments to be passed to compression program.</p>
+      </td>
+    </tr>
     <tr>
       <td><code>mode</code></td>
       <td>

pkg/pkg.bzl

@@ -76,6 +76,8 @@ def _pkg_tar_impl(ctx):
         "--owner=" + ctx.attr.owner,
         "--owner_name=" + ctx.attr.ownername,
     ]
+    if ctx.executable.compressor:
+        args.append("--compressor=%s %s" % (ctx.executable.compressor.path, ctx.attr.compressor_args))
     if ctx.attr.mtime != _DEFAULT_MTIME:
         if ctx.attr.portable_mtime:
             fail("You may not set both mtime and portable_mtime")
@@ -129,7 +131,7 @@ def _pkg_tar_impl(ctx):
         args += ["--empty_file=%s" % empty_file for empty_file in ctx.attr.empty_files]
     if ctx.attr.empty_dirs:
         args += ["--empty_dir=%s" % empty_dir for empty_dir in ctx.attr.empty_dirs]
-    if ctx.attr.extension:
+    if ctx.attr.extension and not ctx.executable.compressor:
         dotPos = ctx.attr.extension.find(".")
         if dotPos > 0:
             dotPos += 1
@@ -149,6 +151,7 @@ def _pkg_tar_impl(ctx):
         mnemonic = "PackageTar",
         progress_message = "Writing: %s" % output_file.path,
         inputs = file_inputs + ctx.files.deps + files,
+        tools = [ctx.executable.compressor],
         executable = ctx.executable.build_tar,
         arguments = ["@" + arg_file.path],
         outputs = [output_file],
@@ -351,6 +354,8 @@ pkg_tar_impl = rule(
         "include_runfiles": attr.bool(),
         "empty_dirs": attr.string_list(),
         "remap_paths": attr.string_dict(),
+        "compressor": attr.label(executable = True, cfg = "exec"),
+        "compressor_args": attr.string(),
 
         # Common attributes
         "out": attr.output(mandatory = True),