Merge branch 'becky/tar_compression' into 'master'

Allowing `bst artifact checkout --tar` to use compression

See merge request BuildStream/buildstream!1451

3 files changed, 75 insertions, 19 deletions

diff --git a/src/buildstream/_frontend/cli.py b/src/buildstream/_frontend/cli.py
index d4686e4fc..bf92161bf 100644
--- a/src/buildstream/_frontend/cli.py
+++ b/src/buildstream/_frontend/cli.py
@@ -10,6 +10,7 @@ from .. import _yaml
 from .._exceptions import BstError, LoadError, AppError
 from .._versions import BST_FORMAT_VERSION
 from .complete import main_bashcomplete, complete_path, CompleteUnhandled
+from ..utils import _get_compression, UtilError
 
 
 ##################################################################
@@ -970,6 +971,9 @@ def artifact():
               help="Create a tarball from the artifact contents instead "
                    "of a file tree. If LOCATION is '-', the tarball "
                    "will be dumped to the standard output.")
+@click.option('--compression', default=None,
+              type=click.Choice(['gz', 'xz', 'bz2']),
+              help="The compression option of the tarball created.")
 @click.option('--pull', 'pull_', default=False, is_flag=True,
               help="Whether to pull the artifact if it's missing or "
                    "incomplete.")
@@ -979,7 +983,7 @@ def artifact():
 @click.argument('element', required=False,
                 type=click.Path(readable=False))
 @click.pass_obj
-def artifact_checkout(app, force, deps, integrate, hardlinks, tar, pull_, directory, element):
+def artifact_checkout(app, force, deps, integrate, hardlinks, tar, compression, pull_, directory, element):
     """Checkout contents of an artifact
 
     When this command is executed from a workspace directory, the default
@@ -987,25 +991,38 @@ def artifact_checkout(app, force, deps, integrate, hardlinks, tar, pull_, direct
     """
     from ..element import Scope
 
-    if hardlinks and tar is not None:
+    if hardlinks and tar:
         click.echo("ERROR: options --hardlinks and --tar conflict", err=True)
         sys.exit(-1)
 
-    if tar is not None and directory is not None:
+    if tar and directory:
         click.echo("ERROR: options --directory and --tar conflict", err=True)
         sys.exit(-1)
 
-    if tar is not None:
-        location = tar
-        tar = True
-    else:
-        if directory is None:
-            location = os.path.abspath(os.path.join(os.getcwd(), element))
+    if not tar:
+        if compression:
+            click.echo("ERROR: --compression can only be provided if --tar is provided", err=True)
+            sys.exit(-1)
         else:
-            location = directory
-        if location[-4:] == '.bst':
-            location = location[:-4]
-        tar = False
+            if directory is None:
+                location = os.path.abspath(os.path.join(os.getcwd(), element))
+            else:
+                location = directory
+            if location[-4:] == '.bst':
+                location = location[:-4]
+            tar = False
+    else:
+        location = tar
+        try:
+            inferred_compression = _get_compression(tar)
+        except UtilError as e:
+            click.echo("ERROR: Invalid file extension given with '--tar': {}".format(e), err=True)
+            sys.exit(-1)
+        if compression and inferred_compression != '' and inferred_compression != compression:
+            click.echo("WARNING: File extension and compression differ."
+                       "File extension has been overridden by --compression", err=True)
+        if not compression:
+            compression = inferred_compression
 
     if deps == "build":
         scope = Scope.BUILD
@@ -1026,8 +1043,9 @@ def artifact_checkout(app, force, deps, integrate, hardlinks, tar, pull_, direct
                             scope=scope,
                             integrate=True if integrate is None else integrate,
                             hardlinks=hardlinks,
-                            tar=tar,
-                            pull=pull_)
+                            pull=pull_,
+                            compression=compression,
+                            tar=bool(tar))
 
 
 ################################################################
diff --git a/src/buildstream/_stream.py b/src/buildstream/_stream.py
index 3c32ff616..bc0f55009 100644
--- a/src/buildstream/_stream.py
+++ b/src/buildstream/_stream.py
@@ -512,8 +512,9 @@ class Stream():
                  scope=Scope.RUN,
                  integrate=True,
                  hardlinks=False,
-                 tar=False,
-                 pull=False):
+                 compression='',
+                 pull=False,
+                 tar=False):
 
         # if pulling we need to ensure dependency artifacts are also pulled
         selection = PipelineSelection.RUN if pull else PipelineSelection.NONE
@@ -552,21 +553,23 @@ class Stream():
                                               .format(e)) from e
                 else:
                     if location == '-':
+                        mode = 'w|' + compression
                         with target.timed_activity("Creating tarball"):
                             # Save the stdout FD to restore later
                             saved_fd = os.dup(sys.stdout.fileno())
                             try:
                                 with os.fdopen(sys.stdout.fileno(), 'wb') as fo:
-                                    with tarfile.open(fileobj=fo, mode="w|") as tf:
+                                    with tarfile.open(fileobj=fo, mode=mode) as tf:
                                         sandbox_vroot.export_to_tar(tf, '.')
                             finally:
                                 # No matter what, restore stdout for further use
                                 os.dup2(saved_fd, sys.stdout.fileno())
                                 os.close(saved_fd)
                     else:
+                        mode = 'w:' + compression
                         with target.timed_activity("Creating tarball '{}'"
                                                    .format(location)):
-                            with tarfile.open(location, "w:") as tf:
+                            with tarfile.open(location, mode=mode) as tf:
                                 sandbox_vroot.export_to_tar(tf, '.')
 
         except BstError as e:
diff --git a/src/buildstream/utils.py b/src/buildstream/utils.py
index 03ff67bf2..82cd4134b 100644
--- a/src/buildstream/utils.py
+++ b/src/buildstream/utils.py
@@ -1351,3 +1351,38 @@ def _deterministic_umask():
         yield
     finally:
         os.umask(old_umask)
+
+
+# _get_compression:
+#
+# Given a file name infer the compression
+#
+# Args:
+#    tar (str): The file name from which to determine compression
+#
+# Returns:
+#    (str): One from '', 'gz', 'xz', 'bz2'
+#
+# Raises:
+#    UtilError: In the case where an unsupported file extension has been provided,
+#               expecting compression.
+#
+#
+def _get_compression(tar):
+    mapped_extensions = {'.tar': '', '.gz': 'gz', '.xz': 'xz', '.bz2': 'bz2'}
+
+    name, ext = os.path.splitext(tar)
+
+    try:
+        return mapped_extensions[ext]
+    except KeyError:
+        # If ext not in mapped_extensions, find out if inner ext is .tar
+        # If so, we assume we have been given an unsupported extension,
+        # which expects compression. Raise an error
+        _, suffix = os.path.splitext(name)
+        if suffix == '.tar':
+            raise UtilError("Expected compression with unknown file extension ('{}'), "
+                            "supported extensions are ('.tar'), ('.gz'), ('.xz'), ('.bz2')".format(ext))
+        else:
+            # Assume just an unconventional name was provided, default to uncompressed
+            return ''