cli.py: Add --compression option to source checkout

!1451 introduced the --compression option to bst artifact checkout.
This MR is a step towards making the commands more symmetric,
and introduces the --compression option to bst source checkout

Now we're compressing, we must explicitly close the tarball
once we're done writing to it.

2 files changed, 20 insertions, 8 deletions

diff --git a/src/buildstream/_frontend/cli.py b/src/buildstream/_frontend/cli.py
index 9fa482666..d657ec76d 100644
--- a/src/buildstream/_frontend/cli.py
+++ b/src/buildstream/_frontend/cli.py
@@ -811,6 +811,9 @@ def source_track(app, elements, deps, except_, cross_junctions):
               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),
@@ -818,7 +821,7 @@ def source_track(app, elements, deps, except_, cross_junctions):
 @click.argument('element', required=False, type=click.Path(readable=False))
 @click.pass_obj
 def source_checkout(app, element, directory, force, deps, except_,
-                    tar, build_scripts):
+                    tar, compression, build_scripts):
     """Checkout sources of an element to the specified location
 
     When this command is executed from a workspace directory, the default
@@ -829,6 +832,10 @@ def source_checkout(app, element, directory, force, deps, except_,
         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)
+
     # 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
@@ -845,6 +852,7 @@ def source_checkout(app, element, directory, force, deps, except_,
                                    deps=deps,
                                    except_targets=except_,
                                    tar=bool(tar),
+                                   compression=compression,
                                    include_build_scripts=build_scripts)
 
 
diff --git a/src/buildstream/_stream.py b/src/buildstream/_stream.py
index 288940c7e..46d81d0f2 100644
--- a/src/buildstream/_stream.py
+++ b/src/buildstream/_stream.py
@@ -741,6 +741,7 @@ class Stream():
     #    deps (str): The dependencies to checkout
     #    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, *,
@@ -749,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)
@@ -764,7 +766,7 @@ 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
@@ -1469,6 +1471,7 @@ class Stream():
                          force=False,
                          deps='none',
                          tar=False,
+                         compression=None,
                          include_build_scripts=False):
         location = os.path.abspath(location)
 
@@ -1481,7 +1484,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:
@@ -1526,16 +1529,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