Merge branch 'jennis/update_source_checkout' into 'master'

Ensure `source checkout` is symmetric to `artifact checkout`

See merge request BuildStream/buildstream!1590

5 files changed, 160 insertions, 49 deletions

diff --git a/NEWS b/NEWS
index f4f85479c..67be6ed98 100644
--- a/NEWS
+++ b/NEWS
@@ -32,11 +32,14 @@ CLI
   o BREAKING CHANGE: `bst init` no longer uses the `--directory` or `-C`
     option. Instead, it (optionally) takes a directory as an argument.
 
-  o BREAKING CHANGE: The bst source-bundle command has been removed. The
+  o BREAKING CHANGE: The `bst source-bundle` command has been removed. The
     functionality it provided has been replaced by the `--include-build-scripts`
-    option of the `bst source-checkout` command. To produce a tarball containing
+    option of the `bst source checkout` command. To produce a tarball containing
     an element's sources and generated build scripts you can do the command
-    `bst source-checkout --include-build-scripts --tar foo.bst some-file.tar`
+    `bst source checkout --include-build-scripts --tar foo.tar foo.bst`.
+
+    A `--compression` option is also supported when using `--tar` which supports
+    xz, gz and bz2 compression.
 
   o BREAKING CHANGE: The 'auto-init' functionality has been removed. This would
     offer to create a project in the event that bst was run against a directory
diff --git a/src/buildstream/_frontend/cli.py b/src/buildstream/_frontend/cli.py
index 3344c5c79..d657ec76d 100644
--- a/src/buildstream/_frontend/cli.py
+++ b/src/buildstream/_frontend/cli.py
@@ -798,7 +798,7 @@ def source_track(app, elements, deps, except_, cross_junctions):
 ##################################################################
 #                  Source Checkout Command                      #
 ##################################################################
-@source.command(name='checkout', short_help='Checkout sources for an element')
+@source.command(name='checkout', short_help='Checkout sources of an element')
 @click.option('--force', '-f', is_flag=True,
               help="Allow files to be overwritten")
 @click.option('--except', 'except_', multiple=True,
@@ -807,28 +807,38 @@ def source_track(app, elements, deps, except_, cross_junctions):
 @click.option('--deps', '-d', default='none', show_default=True,
               type=click.Choice(['build', 'none', 'run', 'all']),
               help='The dependencies whose sources to checkout')
-@click.option('--tar', 'tar', is_flag=True,
-              help='Create a tarball from the element\'s sources instead of a '
-                   'file tree.')
+@click.option('--tar', default=None, metavar='LOCATION',
+              type=click.Path(),
+              help="Create a tarball containing the sources instead "
+                   "of a file tree.")
+@click.option('--compression', default=None,
+              type=click.Choice(['gz', 'xz', 'bz2']),
+              help="The compression option of the tarball created.")
 @click.option('--include-build-scripts', 'build_scripts', is_flag=True)
+@click.option('--directory', default='source-checkout',
+              type=click.Path(file_okay=False),
+              help="The directory to checkout the sources to")
 @click.argument('element', required=False, type=click.Path(readable=False))
-@click.argument('location', type=click.Path(), required=False)
 @click.pass_obj
-def source_checkout(app, element, location, force, deps, except_,
-                    tar, build_scripts):
+def source_checkout(app, element, directory, force, deps, except_,
+                    tar, compression, build_scripts):
     """Checkout sources of an element to the specified location
 
     When this command is executed from a workspace directory, the default
     is to checkout the sources of the workspace element.
     """
-    if not element and not location:
-        click.echo("ERROR: LOCATION is not specified", err=True)
+
+    if tar and directory != "source-checkout":
+        click.echo("ERROR: options --directory and --tar conflict", err=True)
+        sys.exit(-1)
+
+    if compression and not tar:
+        click.echo("ERROR: --compression specified without --tar", err=True)
         sys.exit(-1)
 
-    if element and not location:
-        # Nasty hack to get around click's optional args
-        location = element
-        element = None
+    # Set the location depending on whether --tar/--directory were specified
+    # Note that if unset, --directory defaults to "source-checkout"
+    location = tar if tar else directory
 
     with app.initialized():
         if not element:
@@ -841,7 +851,8 @@ def source_checkout(app, element, location, force, deps, except_,
                                    force=force,
                                    deps=deps,
                                    except_targets=except_,
-                                   tar=tar,
+                                   tar=bool(tar),
+                                   compression=compression,
                                    include_build_scripts=build_scripts)
 
 
diff --git a/src/buildstream/_stream.py b/src/buildstream/_stream.py
index 293ba051d..a637ef12d 100644
--- a/src/buildstream/_stream.py
+++ b/src/buildstream/_stream.py
@@ -572,6 +572,18 @@ class Stream():
             raise StreamError("Error while staging dependencies into a sandbox"
                               ": '{}'".format(e), detail=e.detail, reason=e.reason) from e
 
+    # _export_artifact()
+    #
+    # Export the files of the artifact/a tarball to a virtual directory
+    #
+    # Args:
+    #    tar (bool): Whether we want to create a tarfile
+    #    location (str): The name of the directory/the tarfile we want to export to/create
+    #    compression (str): The type of compression for the tarball
+    #    target (Element/ArtifactElement): The Element/ArtifactElement we want to checkout
+    #    hardlinks (bool): Whether to checkout hardlinks instead of copying
+    #    virdir (Directory): The sandbox's root directory as a virtual directory
+    #
     def _export_artifact(self, tar, location, compression, target, hardlinks, virdir):
         if not tar:
             with target.timed_activity("Checking out files in '{}'"
@@ -585,9 +597,10 @@ class Stream():
                     raise StreamError("Failed to checkout files: '{}'"
                                       .format(e)) from e
         else:
-            if location == '-':
-                mode = 'w|' + compression
-                with target.timed_activity("Creating tarball"):
+            to_stdout = location == '-'
+            mode = _handle_compression(compression, to_stream=to_stdout)
+            with target.timed_activity("Creating tarball"):
+                if to_stdout:
                     # Save the stdout FD to restore later
                     saved_fd = os.dup(sys.stdout.fileno())
                     try:
@@ -598,10 +611,7 @@ class Stream():
                         # 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)):
+                else:
                     with tarfile.open(location, mode=mode) as tf:
                         virdir.export_to_tar(tf, '.')
 
@@ -727,8 +737,12 @@ class Stream():
     # Args:
     #    target (str): The target element whose sources to checkout
     #    location (str): Location to checkout the sources to
+    #    force (bool): Whether to overwrite existing directories/tarfiles
     #    deps (str): The dependencies to checkout
-    #    except_targets (list): List of targets to except from staging
+    #    except_targets ([str]): List of targets to except from staging
+    #    tar (bool): Whether to write a tarfile holding the checkout contents
+    #    compression (str): The type of compression for tarball
+    #    include_build_scripts (bool): Whether to include build scripts in the checkout
     #
     def source_checkout(self, target, *,
                         location=None,
@@ -736,6 +750,7 @@ class Stream():
                         deps='none',
                         except_targets=(),
                         tar=False,
+                        compression=None,
                         include_build_scripts=False):
 
         self._check_location_writable(location, force=force, tar=tar)
@@ -751,11 +766,13 @@ class Stream():
         # Stage all sources determined by scope
         try:
             self._source_checkout(elements, location, force, deps,
-                                  tar, include_build_scripts)
+                                  tar, compression, include_build_scripts)
         except BstError as e:
             raise StreamError("Error while writing sources"
                               ": '{}'".format(e), detail=e.detail, reason=e.reason) from e
 
+        self._message(MessageType.INFO, "Checked out sources to '{}'".format(location))
+
     # workspace_open
     #
     # Open a project workspace
@@ -1456,6 +1473,7 @@ class Stream():
                          force=False,
                          deps='none',
                          tar=False,
+                         compression=None,
                          include_build_scripts=False):
         location = os.path.abspath(location)
 
@@ -1468,7 +1486,7 @@ class Stream():
             if include_build_scripts:
                 self._write_build_scripts(temp_source_dir.name, elements)
             if tar:
-                self._create_tarball(temp_source_dir.name, location)
+                self._create_tarball(temp_source_dir.name, location, compression)
             else:
                 self._move_directory(temp_source_dir.name, location, force)
         except OSError as e:
@@ -1513,16 +1531,17 @@ class Stream():
                 element._stage_sources_at(element_source_dir, mount_workspaces=False)
 
     # Create a tarball from the content of directory
-    def _create_tarball(self, directory, tar_name):
+    def _create_tarball(self, directory, tar_name, compression):
+        if compression is None:
+            compression = ''
+        mode = _handle_compression(compression)
         try:
             with utils.save_file_atomic(tar_name, mode='wb') as f:
-                # This TarFile does not need to be explicitly closed
-                # as the underlying file object will be closed be the
-                # save_file_atomic contect manager
-                tarball = tarfile.open(fileobj=f, mode='w')
+                tarball = tarfile.open(fileobj=f, mode=mode)
                 for item in os.listdir(str(directory)):
                     file_to_add = os.path.join(directory, item)
                     tarball.add(file_to_add, arcname=item)
+                tarball.close()
         except OSError as e:
             raise StreamError("Failed to create tar archive: {}".format(e)) from e
 
@@ -1696,3 +1715,19 @@ class Stream():
         # a new use-case arises.
         #
         raise TypeError("Stream objects should not be pickled.")
+
+
+# _handle_compression()
+#
+# Return the tarfile mode str to be used when creating a tarball
+#
+# Args:
+#    compression (str): The type of compression (either 'gz', 'xz' or 'bz2')
+#    to_stdout (bool): Whether we want to open a stream for writing
+#
+# Returns:
+#    (str): The tarfile mode string
+#
+def _handle_compression(compression, *, to_stream=False):
+    mode_prefix = 'w|' if to_stream else 'w:'
+    return mode_prefix + compression
diff --git a/tests/frontend/source_checkout.py b/tests/frontend/source_checkout.py
index 1831ee863..8d6bae83b 100644
--- a/tests/frontend/source_checkout.py
+++ b/tests/frontend/source_checkout.py
@@ -52,7 +52,7 @@ def test_source_checkout(datafiles, cli, tmpdir_factory, with_workspace, guess_e
     else:
         ws_cmd = []
 
-    args = ws_cmd + ['source', 'checkout', '--deps', 'none', *elm_cmd, checkout]
+    args = ws_cmd + ['source', 'checkout', '--deps', 'none', '--directory', checkout, *elm_cmd]
     result = cli.run(project=project, args=args)
     result.assert_success()
 
@@ -66,10 +66,14 @@ def test_source_checkout_force(datafiles, cli, force_flag):
     checkout = os.path.join(cli.directory, 'source-checkout')
     target = 'checkout-deps.bst'
 
+    # Make the checkout directory with 'some-thing' inside it
     os.makedirs(os.path.join(checkout, 'some-thing'))
-    # Path(os.path.join(checkout, 'some-file')).touch()
 
-    result = cli.run(project=project, args=['source', 'checkout', force_flag, target, '--deps', 'none', checkout])
+    result = cli.run(project=project, args=['source', 'checkout',
+                                            force_flag,
+                                            '--deps', 'none',
+                                            '--directory', checkout,
+                                            target])
     result.assert_success()
 
     assert os.path.exists(os.path.join(checkout, 'checkout-deps', 'etc', 'buildstream', 'config'))
@@ -78,28 +82,52 @@ def test_source_checkout_force(datafiles, cli, force_flag):
 @pytest.mark.datafiles(DATA_DIR)
 def test_source_checkout_tar(datafiles, cli):
     project = str(datafiles)
-    checkout = os.path.join(cli.directory, 'source-checkout.tar')
+    tar = os.path.join(cli.directory, 'source-checkout.tar')
     target = 'checkout-deps.bst'
 
-    result = cli.run(project=project, args=['source', 'checkout', '--tar', target, '--deps', 'none', checkout])
+    result = cli.run(project=project, args=['source', 'checkout',
+                                            '--tar', tar,
+                                            '--deps', 'none',
+                                            target])
     result.assert_success()
 
-    assert os.path.exists(checkout)
-    with tarfile.open(checkout) as tf:
-        expected_content = os.path.join(checkout, 'checkout-deps', 'etc', 'buildstream', 'config')
+    assert os.path.exists(tar)
+    with tarfile.open(tar) as tf:
+        expected_content = os.path.join(tar, 'checkout-deps', 'etc', 'buildstream', 'config')
         tar_members = [f.name for f in tf]
         for member in tar_members:
             assert member in expected_content
 
 
 @pytest.mark.datafiles(DATA_DIR)
+@pytest.mark.parametrize("compression", [("gz"), ("xz"), ("bz2")])
+def test_source_checkout_compressed_tar(datafiles, cli, compression):
+    project = str(datafiles)
+    tarfile_name = "source-checkout.tar" + compression
+    tar = os.path.join(cli.directory, tarfile_name)
+    target = 'checkout-deps.bst'
+
+    result = cli.run(project=project, args=['source', 'checkout',
+                                            '--tar', tar,
+                                            '--compression', compression,
+                                            '--deps', 'none',
+                                            target])
+    result.assert_success()
+    tar = tarfile.open(name=tar, mode='r:' + compression)
+    assert os.path.join('checkout-deps', 'etc', 'buildstream', 'config') in tar.getnames()
+
+
+@pytest.mark.datafiles(DATA_DIR)
 @pytest.mark.parametrize('deps', [('build'), ('none'), ('run'), ('all')])
 def test_source_checkout_deps(datafiles, cli, deps):
     project = str(datafiles)
     checkout = os.path.join(cli.directory, 'source-checkout')
     target = 'checkout-deps.bst'
 
-    result = cli.run(project=project, args=['source', 'checkout', target, '--deps', deps, checkout])
+    result = cli.run(project=project, args=['source', 'checkout',
+                                            '--directory', checkout,
+                                            '--deps', deps,
+                                            target])
     result.assert_success()
 
     # Sources of the target
@@ -127,10 +155,11 @@ def test_source_checkout_except(datafiles, cli):
     checkout = os.path.join(cli.directory, 'source-checkout')
     target = 'checkout-deps.bst'
 
-    result = cli.run(project=project, args=['source', 'checkout', target,
+    result = cli.run(project=project, args=['source', 'checkout',
+                                            '--directory', checkout,
                                             '--deps', 'all',
                                             '--except', 'import-bin.bst',
-                                            checkout])
+                                            target])
     result.assert_success()
 
     # Sources for the target should be present
@@ -162,7 +191,7 @@ def test_source_checkout_fetch(datafiles, cli):
 
     args = ['source', 'checkout']
     args += [target, checkout]
-    result = cli.run(project=project, args=args)
+    result = cli.run(project=project, args=['source', 'checkout', '--directory', checkout, target])
 
     result.assert_success()
     assert os.path.exists(os.path.join(checkout, 'remote-import-dev', 'pony.h'))
@@ -175,7 +204,7 @@ def test_source_checkout_build_scripts(cli, tmpdir, datafiles):
     normal_name = 'source-bundle-source-bundle-hello'
     checkout = os.path.join(str(tmpdir), 'source-checkout')
 
-    args = ['source', 'checkout', '--include-build-scripts', element_name, checkout]
+    args = ['source', 'checkout', '--include-build-scripts', '--directory', checkout, element_name]
     result = cli.run(project=project_path, args=args)
     result.assert_success()
 
@@ -192,7 +221,7 @@ def test_source_checkout_tar_buildscripts(cli, tmpdir, datafiles):
     normal_name = 'source-bundle-source-bundle-hello'
     tar_file = os.path.join(str(tmpdir), 'source-checkout.tar')
 
-    args = ['source', 'checkout', '--include-build-scripts', '--tar', element_name, tar_file]
+    args = ['source', 'checkout', '--include-build-scripts', '--tar', tar_file, element_name]
     result = cli.run(project=project_path, args=args)
     result.assert_success()
 
@@ -201,3 +230,36 @@ def test_source_checkout_tar_buildscripts(cli, tmpdir, datafiles):
     with tarfile.open(tar_file, 'r') as tf:
         for script in expected_scripts:
             assert script in tf.getnames()
+
+
+# Test that the --directory and --tar options conflict
+@pytest.mark.datafiles(DATA_DIR)
+def test_source_checkout_options_tar_and_dir_conflict(cli, tmpdir, datafiles):
+    project = str(datafiles)
+    checkout = os.path.join(cli.directory, 'source-checkout')
+    tar_file = os.path.join(str(tmpdir), 'source-checkout.tar')
+    target = 'checkout-deps.bst'
+
+    result = cli.run(project=project, args=['source', 'checkout',
+                                            '--directory', checkout,
+                                            '--tar', tar_file,
+                                            target])
+
+    assert result.exit_code != 0
+    assert "ERROR: options --directory and --tar conflict" in result.stderr
+
+
+# Test that the --compression option without --tar fails
+@pytest.mark.datafiles(DATA_DIR)
+def test_source_checkout_compression_without_tar(cli, tmpdir, datafiles):
+    project = str(datafiles)
+    checkout = os.path.join(cli.directory, 'source-checkout')
+    target = 'checkout-deps.bst'
+
+    result = cli.run(project=project, args=['source', 'checkout',
+                                            '--directory', checkout,
+                                            '--compression', 'xz',
+                                            target])
+
+    assert result.exit_code != 0
+    assert "ERROR: --compression specified without --tar" in result.stderr
diff --git a/tests/sourcecache/source-checkout.py b/tests/sourcecache/source-checkout.py
index c2c7fe3cd..4e3391c12 100644
--- a/tests/sourcecache/source-checkout.py
+++ b/tests/sourcecache/source-checkout.py
@@ -49,7 +49,7 @@ def test_source_checkout(tmpdir, datafiles, cli):
     repo = create_element_size('target.bst', project_dir, element_path, [], 100000)
 
     # check implicit fetching
-    res = cli.run(project=project_dir, args=['source', 'checkout', 'target.bst', target_dir])
+    res = cli.run(project=project_dir, args=['source', 'checkout', '--directory', target_dir, 'target.bst'])
     res.assert_success()
     assert "Fetching from" in res.stderr
 
@@ -60,7 +60,7 @@ def test_source_checkout(tmpdir, datafiles, cli):
     shutil.rmtree(source_dir)
 
     res = cli.run(project=project_dir,
-                  args=['source', 'checkout', 'target.bst', target_dir])
+                  args=['source', 'checkout', '--directory', target_dir, 'target.bst'])
     res.assert_success()
     assert "Fetching from" not in res.stderr
 
@@ -69,5 +69,5 @@ def test_source_checkout(tmpdir, datafiles, cli):
     shutil.rmtree(os.path.join(cache_dir, 'cas'))
 
     res = cli.run(project=project_dir,
-                  args=['source', 'checkout', 'target.bst', target_dir])
+                  args=['source', 'checkout', '--directory', target_dir, 'target.bst'])
     res.assert_task_error(ErrorDomain.PLUGIN, None)