WIP Add support for dumping a tarball streamtiagogomes/tarball

2 files changed, 69 insertions, 26 deletions

diff --git a/buildstream/_frontend/cli.py b/buildstream/_frontend/cli.py
index e59b1baec..10c7c2804 100644
--- a/buildstream/_frontend/cli.py
+++ b/buildstream/_frontend/cli.py
@@ -627,19 +627,22 @@ def shell(app, element, sysroot, mount, isolate, build_, command):
               help="Whether to run integration commands")
 @click.option('--hardlinks', default=False, is_flag=True,
               help="Checkout hardlinks instead of copies (handle with care)")
+@click.option('--tar', default=False, is_flag=True,
+              help="Create a tarball of the artifact contents")
 @click.argument('element',
                 type=click.Path(readable=False))
-@click.argument('directory', type=click.Path(file_okay=False))
+@click.argument('location', type=click.Path())
 @click.pass_obj
-def checkout(app, element, directory, force, integrate, hardlinks):
+def checkout(app, element, location, force, integrate, hardlinks, tar):
     """Checkout a built artifact to the specified directory
     """
     with app.initialized():
         app.stream.checkout(element,
-                            directory=directory,
+                            location=location,
                             force=force,
                             integrate=integrate,
-                            hardlinks=hardlinks)
+                            hardlinks=hardlinks,
+                            tar=tar)
 
 
 ##################################################################
diff --git a/buildstream/_stream.py b/buildstream/_stream.py
index 5013daf3b..8b84b6c33 100644
--- a/buildstream/_stream.py
+++ b/buildstream/_stream.py
@@ -20,6 +20,7 @@
 #        Tristan Maat <tristan.maat@codethink.co.uk>
 
 import os
+import sys
 import stat
 import shlex
 import shutil
@@ -359,28 +360,37 @@ class Stream():
     #    integrate (bool): Whether to run integration commands
     #    hardlinks (bool): Whether checking out files hardlinked to
     #                      their artifacts is acceptable
+    #    tar (bool): Create a tarball of the artifact contents
     #
     def checkout(self, target, *,
-                 directory=None,
+                 location=None,
                  force=False,
                  integrate=True,
-                 hardlinks=False):
+                 hardlinks=False,
+                 tar=False):
 
         # We only have one target in a checkout command
         elements, _ = self._load((target,), (), fetch_subprojects=True)
         target = elements[0]
 
-        try:
-            os.makedirs(directory, exist_ok=True)
-        except OSError as e:
-            raise StreamError("Failed to create checkout directory: {}".format(e)) from e
-
-        if not os.access(directory, os.W_OK):
-            raise StreamError("Directory {} not writable".format(directory))
-
-        if not force and os.listdir(directory):
-            raise StreamError("Checkout directory is not empty: {}"
-                              .format(directory))
+        if not tar:
+            try:
+                os.makedirs(location, exist_ok=True)
+            except OSError as e:
+                raise StreamError("Failed to create checkout directory: '{}'".format(e)) from e
+
+        if not tar:
+            if not os.access(location, os.W_OK):
+                raise StreamError("Checkout directory '{}' not writable".format(location))
+            if not force and os.listdir(location):
+                raise StreamError("Checkout directory '{}' not empty"
+                                  .format(location))
+        elif os.path.exists(location) and location != '-':
+            if not os.access(location, os.W_OK):
+                raise StreamError("Output file '{}' not writable".format(location))
+            if not force and os.path.exists(location):
+                raise StreamError("Output file '{}' already exists"
+                                  .format(location))
 
         # Stage deps into a temporary sandbox first
         try:
@@ -388,16 +398,28 @@ class Stream():
 
                 # Copy or move the sandbox to the target directory
                 sandbox_root = sandbox.get_directory()
-                with target.timed_activity("Checking out files in {}".format(directory)):
-                    try:
-                        if hardlinks:
-                            self._checkout_hardlinks(sandbox_root, directory)
-                        else:
-                            utils.copy_files(sandbox_root, directory)
-                    except OSError as e:
-                        raise StreamError("Failed to checkout files: {}".format(e)) from e
+                if not tar:
+                    with target.timed_activity("Checking out files in '{}'".format(location)):
+                        try:
+                            if hardlinks:
+                                self._checkout_hardlinks(sandbox_root, location)
+                            else:
+                                utils.copy_files(sandbox_root, location)
+                        except OSError as e:
+                            raise StreamError("Failed to checkout files: '{}'".format(e)) from e
+                else:
+                    if location != '-':
+                        with target.timed_activity("Creating tarball '{}'".format(location)):
+                            with tarfile.open(location, "w:") as tar:
+                                Stream._add_directory_to_tarfile(tar, sandbox_root, '.')
+                    else:
+                        with target.timed_activity("Creating tarball"):
+                            with os.fdopen(sys.stdout.fileno(), 'wb') as fo:
+                                with tarfile.open(fileobj=fo, mode="w|") as tar:
+                                    Stream._add_directory_to_tarfile(tar, sandbox_root, '.')
+
         except BstError as e:
-            raise StreamError("Error while staging dependencies into a sandbox: {}".format(e),
+            raise StreamError("Error while staging dependencies into a sandbox: '{}'".format(e),
                               reason=e.reason) from e
 
     # workspace_open
@@ -1007,6 +1029,24 @@ class Stream():
         else:
             utils.link_files(sandbox_root, directory)
 
+    @staticmethod
+    def _add_directory_to_tarfile(tar, dir_name, dir_arcname):
+        for filename in sorted(os.listdir(dir_name)):
+            name = os.path.join(dir_name, filename)
+            arcname = os.path.join(dir_arcname, filename)
+
+            tarinfo = tar.gettarinfo(name, arcname)
+            tarinfo.mtime = 0
+
+            if tarinfo.isreg():
+                with open(name, "rb") as f:
+                    tar.addfile(tarinfo, f)
+            elif tarinfo.isdir():
+                tar.addfile(tarinfo)
+                Stream._add_directory_to_tarfile(tar, name, arcname)
+            else:
+                tar.addfile(tarinfo)
+
     # Write the element build script to the given directory
     def _write_element_script(self, directory, element):
         try: