Allowing tar files to be compressedbecky/tar_compression

Changes made to cli.py and _stream.py in order to support tar
compression. Compression flag has been added, which overrides any file
extensions given. Where no compression or file extension provided,
default to uncompressed .tar.

4 files changed, 168 insertions, 21 deletions

diff --git a/src/buildstream/_frontend/cli.py b/src/buildstream/_frontend/cli.py
index 3af4121ea..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
@@ -991,25 +995,34 @@ def artifact_checkout(app, force, deps, integrate, hardlinks, tar, pull_, direct
         click.echo("ERROR: options --hardlinks and --tar conflict", err=True)
         sys.exit(-1)
 
-    if not tar and not directory:
-        click.echo("ERROR: One of --directory or --tar must be provided", err=True)
-        sys.exit(-1)
-
     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
@@ -1030,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 ''
diff --git a/tests/frontend/buildcheckout.py b/tests/frontend/buildcheckout.py
index aadefc2ab..a3f68c031 100644
--- a/tests/frontend/buildcheckout.py
+++ b/tests/frontend/buildcheckout.py
@@ -211,6 +211,76 @@ def test_build_checkout_unbuilt(datafiles, cli):
 
 
 @pytest.mark.datafiles(DATA_DIR)
+def test_build_checkout_compression_no_tar(datafiles, cli):
+    project = str(datafiles)
+    checkout = os.path.join(cli.directory, 'checkout.tar')
+
+    result = cli.run(project=project, args=['build', 'target.bst'])
+    result.assert_success()
+
+    checkout_args = ['artifact', 'checkout', '--directory', checkout, '--compression', 'gz', 'target.bst']
+
+    result = cli.run(project=project, args=checkout_args)
+    assert "ERROR: --compression can only be provided if --tar is provided" in result.stderr
+    assert result.exit_code != 0
+
+# If we don't support the extension, we default to an uncompressed tarball
+@pytest.mark.datafiles(DATA_DIR)
+def test_build_checkout_tar_with_unconventional_name(datafiles, cli):
+    project = str(datafiles)
+    checkout = os.path.join(cli.directory, 'checkout.foo')
+
+    result = cli.run(project=project, args=['build', 'target.bst'])
+    result.assert_success()
+
+    checkout_args = ['artifact', 'checkout', '--tar', checkout, 'target.bst']
+
+    result = cli.run(project=project, args=checkout_args)
+    result.assert_success()
+
+    tar = tarfile.open(name=checkout, mode='r')
+    assert os.path.join('.', 'usr', 'bin', 'hello') in tar.getnames()
+    assert os.path.join('.', 'usr', 'include', 'pony.h') in tar.getnames()
+
+
+@pytest.mark.datafiles(DATA_DIR)
+def test_build_checkout_tar_with_unsupported_ext(datafiles, cli):
+    project = str(datafiles)
+    checkout = os.path.join(cli.directory, 'checkout.tar.foo')
+
+    result = cli.run(project=project, args=['build', 'target.bst'])
+    result.assert_success()
+
+    checkout_args = ['artifact', 'checkout', '--tar', checkout, 'target.bst']
+
+    result = cli.run(project=project, args=checkout_args)
+    assert "Invalid file extension given with '--tar': Expected compression with unknown file extension ('.foo'), " \
+           "supported extensions are ('.tar'), ('.gz'), ('.xz'), ('.bz2')" in result.stderr
+
+
+@pytest.mark.datafiles(DATA_DIR)
+def test_build_checkout_tar_no_compression(datafiles, cli):
+    project = str(datafiles)
+    checkout = os.path.join(cli.directory, 'checkout.tar.gz')
+
+    result = cli.run(project=project, args=['build', 'target.bst'])
+    result.assert_success()
+
+    builddir = os.path.join(cli.directory, 'build')
+    assert os.path.isdir(builddir)
+    assert not os.listdir(builddir)
+
+    checkout_args = ['artifact', 'checkout', '--tar', checkout, 'target.bst']
+
+    result = cli.run(project=project, args=checkout_args)
+    result.assert_success()
+
+    tar = tarfile.open(name=checkout, mode='r:gz')
+    assert os.path.join('.', 'usr', 'bin', 'hello') in tar.getnames()
+    assert os.path.join('.', 'usr', 'include', 'pony.h') in tar.getnames()
+
+
+@pytest.mark.datafiles(DATA_DIR)
 def test_build_checkout_tarball(datafiles, cli):
     project = str(datafiles)
     checkout = os.path.join(cli.directory, 'checkout.tar')
@@ -251,6 +321,31 @@ def test_build_checkout_no_tar_no_directory(datafiles, cli, tmpdir):
 
 
 @pytest.mark.datafiles(DATA_DIR)
+@pytest.mark.parametrize("compression", [("gz"), ("xz"), ("bz2")])
+def test_build_checkout_tarball_compression(datafiles, cli, compression):
+    project = str(datafiles)
+    checkout = os.path.join(cli.directory, 'checkout.tar')
+
+    result = cli.run(project=project, args=['build', 'target.bst'])
+    result.assert_success()
+
+    builddir = os.path.join(cli.directory, 'build')
+    assert os.path.isdir(builddir)
+    assert not os.listdir(builddir)
+
+    checkout_args = ['artifact', 'checkout', '--tar', checkout, '--compression', compression, 'target.bst']
+
+    result = cli.run(project=project, args=checkout_args)
+    result.assert_success()
+
+    # Docs say not to use TarFile class directly, using .open seems to work.
+    # https://docs.python.org/3/library/tarfile.html#tarfile.TarFile
+    tar = tarfile.open(name=checkout, mode='r:' + compression)
+    assert os.path.join('.', 'usr', 'bin', 'hello') in tar.getnames()
+    assert os.path.join('.', 'usr', 'include', 'pony.h') in tar.getnames()
+
+
+@pytest.mark.datafiles(DATA_DIR)
 def test_build_checkout_tarball_stdout(datafiles, cli):
     project = str(datafiles)
     tarball = os.path.join(cli.directory, 'tarball.tar')