Merge pull request #6 from RyeMutt/zstd

Add zstandard support for tarfile compression along with xz, and gzip

autobuild/autobuild_tool_install.py

@@ -267,7 +267,7 @@ def _install_package(archive_path, install_dir, exclude=[]):
         logger.error("cannot extract non-existing package: %s" % archive_path)
         return False
     logger.info("extracting from %s" % os.path.basename(archive_path))
-    if tarfile.is_tarfile(archive_path):
+    if tarfile.is_tarfile(archive_path) or ".tar.zst" in archive_path:
         return __extract_tar_file(archive_path, install_dir, exclude=exclude)
     elif zipfile.is_zipfile(archive_path):
         return __extract_zip_archive(archive_path, install_dir, exclude=exclude)
@@ -285,8 +285,11 @@ def extract_metadata_from_package(archive_path, metadata_file_name):
         logger.error("no package found at: %s" % archive_path)
     else:
         logger.debug("extracting metadata from %s" % os.path.basename(archive_path))
-        if tarfile.is_tarfile(archive_path):
-            tar = tarfile.open(archive_path, 'r')
+        if tarfile.is_tarfile(archive_path) or ".tar.zst" in archive_path:
+            if ".tar.zst" in archive_path:
+                tar = common.ZstdTarFile(archive_path, 'r')
+            else:
+                tar = tarfile.open(archive_path, 'r')
             try:
                 metadata_file = tar.extractfile(metadata_file_name)
             except KeyError as err:
@@ -305,7 +308,10 @@ def extract_metadata_from_package(archive_path, metadata_file_name):
 
 def __extract_tar_file(cachename, install_dir, exclude=[]):
     # Attempt to extract the package from the install cache
-    tar = tarfile.open(cachename, 'r')
+    if ".tar.zst" in cachename:
+        tar = common.ZstdTarFile(cachename, 'r')
+    else:
+        tar = tarfile.open(cachename, 'r')
     extract = [member for member in tar.getmembers() if member.name not in exclude]
     conflicts = [member.name for member in extract
                  if os.path.exists(os.path.join(install_dir, member.name))

autobuild/autobuild_tool_package.py

@@ -65,7 +65,7 @@ def register(self, parser):
         parser.add_argument('--archive-format',
                             default=None,
                             dest='archive_format',
-                            help='the format of the archive (tbz2 or zip)')
+                            help='the format of the archive (tbz2, tzst, txz, tgz, or zip)')
         parser.add_argument('--build-dir',
                             default=None,
                             dest='select_dir',  # see common.select_directories()
@@ -239,8 +239,8 @@ def package(config, build_directory, platform_name, archive_filename=None, archi
     else:
         archive_description = platform_description.archive
         format = _determine_archive_format(archive_format, archive_description)
-        if format == 'tbz2':
-            _create_tarfile(tarfilename + '.tar.bz2', build_directory, files, results)
+        if format in ('txz', 'tbz2', 'tgz', 'tzst'):
+            _create_tarfile(tarfilename, format, build_directory, files, results)
         elif format == 'zip':
             _create_zip_archive(tarfilename + '.zip', build_directory, files, results)
         else:
@@ -288,13 +288,27 @@ def _get_file_list(platform_description, build_directory):
             os.chdir(current_directory)
     return [files, missing]
 
-def _create_tarfile(tarfilename, build_directory, filelist, results: dict):
+def _create_tarfile(tarfilename, format, build_directory, filelist, results: dict):
     if not os.path.exists(os.path.dirname(tarfilename)):
         os.makedirs(os.path.dirname(tarfilename))
     current_directory = os.getcwd()
     os.chdir(build_directory)
     try:
-        tfile = tarfile.open(tarfilename, 'w:bz2')
+        if format == 'txz':
+            tarfilename = tarfilename + '.tar.xz'
+            tfile = tarfile.open(tarfilename, 'w:xz')
+        elif format == 'tbz2':
+            tarfilename = tarfilename + '.tar.bz2'
+            tfile = tarfile.open(tarfilename, 'w:bz2')
+        elif format == 'tgz':
+            tarfilename = tarfilename + '.tar.gz'
+            tfile = tarfile.open(tarfilename, 'w:gz')
+        elif format == 'tzst':
+            tarfilename = tarfilename + '.tar.zst'
+            tfile = common.ZstdTarFile(tarfilename, 'w', level=22)
+        else:
+            raise PackageError("unknown tar archive format: %s" % format)
+
         for file in filelist:
             try:
                 # Make sure permissions are set on Windows.

autobuild/common.py

@@ -9,14 +9,18 @@
 
 import itertools
 import logging
+import multiprocessing
 import os
 import platform
 import pprint
 import subprocess
 import sys
+import tarfile
 import tempfile
 from collections import OrderedDict
 
+from pyzstd import CParameter, ZstdFile
+
 from autobuild.version import AUTOBUILD_VERSION_STRING
 
 logger = logging.getLogger(__name__)
@@ -535,3 +539,26 @@ def has_cmd(name) -> bool:
     except OSError:
         return False
     return not p.returncode
+
+
+class ZstdTarFile(tarfile.TarFile):
+    def __init__(self, name, mode='r', *, level=4, zstd_dict=None, **kwargs):
+        zstdoption = None
+        if mode != 'r' and mode != 'rb':
+           zstdoption = {CParameter.compressionLevel : level,
+                         CParameter.nbWorkers : multiprocessing.cpu_count(),
+                         CParameter.checksumFlag : 1}
+        self.zstd_file = ZstdFile(name, mode,
+                                level_or_option=zstdoption,
+                                zstd_dict=zstd_dict)
+        try:
+            super().__init__(fileobj=self.zstd_file, mode=mode, **kwargs)
+        except:
+            self.zstd_file.close()
+            raise
+
+    def close(self):
+        try:
+            super().close()
+        finally:
+            self.zstd_file.close()

setup.py

@@ -30,7 +30,7 @@
         'Programming Language :: Python :: 3',
         'Topic :: Software Development :: Libraries :: Python Modules',
     ],
-    install_requires=['llsd', 'pydot'],
+    install_requires=['llsd', 'pydot', 'pyzstd'],
     extras_require={
         'dev': ['pytest', 'pytest-cov'],
         'build': ['build', 'setuptools_scm'],

tests/test_graph.py

@@ -92,7 +92,7 @@ def setUp(self):
         self.options.graph_type = 'mermaid'
 
     def test_output(self):
-        with CaptureStdout() as out: 
+        with CaptureStdout() as out:
             self.options.source_file = os.path.join(self.this_dir, "data", "bongo-0.1-common-111.tar.bz2")
             graph.AutobuildTool().run(self.options)
         graph_txt = out.getvalue()

tests/test_install.py

@@ -510,9 +510,9 @@ def test_versions(self):
         self.assertEqual(stream.getvalue(), "bogus: 1\n")
 
 # -------------------------------------  -------------------------------------
-class TestInstallCachedArchive(BaseTest):
+class TestInstallCachedBZ2Archive(BaseTest):
     def setup_method(self, method):
-        super(TestInstallCachedArchive, self).setup_method(method)
+        super(TestInstallCachedBZ2Archive, self).setup_method(method)
         self.pkg = "bogus"
         # make sure that the archive is not in the server
         clean_file(os.path.join(SERVER_DIR, "bogus-0.1-common-111.tar.bz2"))
@@ -532,6 +532,98 @@ def test_success(self):
             autobuild_tool_install.AutobuildTool().run(self.options)
         self.assertEqual(stream.getvalue(), 'Dirty Packages: \n')
 
+# -------------------------------------  -------------------------------------
+class TestInstallCachedGZArchive(BaseTest):
+    def setup_method(self, method):
+        super(TestInstallCachedGZArchive, self).setup_method(method)
+        self.pkg = "bogus"
+        # make sure that the archive is not in the server
+        clean_file(os.path.join(SERVER_DIR, "bogus-0.1-common-111.tar.gz"))
+        assert not os.path.exists(in_dir(SERVER_DIR, "bogus-0.1-common-111.tar.gz"))
+        self.cache_name = in_dir(common.get_install_cache_dir(), "bogus-0.1-common-111.tar.gz")
+
+        # Instead copy directly to cache dir.
+        self.copyto(os.path.join(mydir, "data", "bogus-0.1-common-111.tar.gz"), common.get_install_cache_dir())
+
+    def test_success(self):
+        self.options.package = [self.pkg]
+        autobuild_tool_install.AutobuildTool().run(self.options)
+        assert os.path.exists(os.path.join(INSTALL_DIR, "lib", "bogus.lib"))
+        assert os.path.exists(os.path.join(INSTALL_DIR, "include", "bogus.h"))
+        with CaptureStdout() as stream:
+            self.options.list_dirty=True
+            autobuild_tool_install.AutobuildTool().run(self.options)
+        self.assertEqual(stream.getvalue(), 'Dirty Packages: \n')
+
+# -------------------------------------  -------------------------------------
+class TestInstallCachedXZArchive(BaseTest):
+    def setup_method(self, method):
+        super(TestInstallCachedXZArchive, self).setup_method(method)
+        self.pkg = "bogus"
+        # make sure that the archive is not in the server
+        clean_file(os.path.join(SERVER_DIR, "bogus-0.1-common-111.tar.xz"))
+        assert not os.path.exists(in_dir(SERVER_DIR, "bogus-0.1-common-111.tar.xz"))
+        self.cache_name = in_dir(common.get_install_cache_dir(), "bogus-0.1-common-111.tar.xz")
+
+        # Instead copy directly to cache dir.
+        self.copyto(os.path.join(mydir, "data", "bogus-0.1-common-111.tar.xz"), common.get_install_cache_dir())
+
+    def test_success(self):
+        self.options.package = [self.pkg]
+        autobuild_tool_install.AutobuildTool().run(self.options)
+        assert os.path.exists(os.path.join(INSTALL_DIR, "lib", "bogus.lib"))
+        assert os.path.exists(os.path.join(INSTALL_DIR, "include", "bogus.h"))
+        with CaptureStdout() as stream:
+            self.options.list_dirty=True
+            autobuild_tool_install.AutobuildTool().run(self.options)
+        self.assertEqual(stream.getvalue(), 'Dirty Packages: \n')
+
+# -------------------------------------  -------------------------------------
+class TestInstallCachedZSTArchive(BaseTest):
+    def setup_method(self, method):
+        super(TestInstallCachedZSTArchive, self).setup_method(method)
+        self.pkg = "bogus"
+        # make sure that the archive is not in the server
+        clean_file(os.path.join(SERVER_DIR, "bogus-0.1-common-111.tar.zst"))
+        assert not os.path.exists(in_dir(SERVER_DIR, "bogus-0.1-common-111.tar.zst"))
+        self.cache_name = in_dir(common.get_install_cache_dir(), "bogus-0.1-common-111.tar.zst")
+
+        # Instead copy directly to cache dir.
+        self.copyto(os.path.join(mydir, "data", "bogus-0.1-common-111.tar.zst"), common.get_install_cache_dir())
+
+    def test_success(self):
+        self.options.package = [self.pkg]
+        autobuild_tool_install.AutobuildTool().run(self.options)
+        assert os.path.exists(os.path.join(INSTALL_DIR, "lib", "bogus.lib"))
+        assert os.path.exists(os.path.join(INSTALL_DIR, "include", "bogus.h"))
+        with CaptureStdout() as stream:
+            self.options.list_dirty=True
+            autobuild_tool_install.AutobuildTool().run(self.options)
+        self.assertEqual(stream.getvalue(), 'Dirty Packages: \n')
+
+# -------------------------------------  -------------------------------------
+class TestInstallCachedZIPArchive(BaseTest):
+    def setup_method(self, method):
+        super(TestInstallCachedZIPArchive, self).setup_method(method)
+        self.pkg = "bogus"
+        # make sure that the archive is not in the server
+        clean_file(os.path.join(SERVER_DIR, "bogus-0.1-common-111.zip"))
+        assert not os.path.exists(in_dir(SERVER_DIR, "bogus-0.1-common-111.zip"))
+        self.cache_name = in_dir(common.get_install_cache_dir(), "bogus-0.1-common-111.zip")
+
+        # Instead copy directly to cache dir.
+        self.copyto(os.path.join(mydir, "data", "bogus-0.1-common-111.zip"), common.get_install_cache_dir())
+
+    def test_success(self):
+        self.options.package = [self.pkg]
+        autobuild_tool_install.AutobuildTool().run(self.options)
+        assert os.path.exists(os.path.join(INSTALL_DIR, "lib", "bogus.lib"))
+        assert os.path.exists(os.path.join(INSTALL_DIR, "include", "bogus.h"))
+        with CaptureStdout() as stream:
+            self.options.list_dirty=True
+            autobuild_tool_install.AutobuildTool().run(self.options)
+        self.assertEqual(stream.getvalue(), 'Dirty Packages: \n')
+
 # -------------------------------------  -------------------------------------
 class TestInstallLocalArchive(BaseTest):
     def setup_method(self, method):

tests/test_package.py

@@ -51,6 +51,9 @@ def setUp(self):
         self.platform = 'common'
         self.tar_basename = os.path.join(self.data_dir, "test1-1.0-common-123456")
         self.tar_name = self.tar_basename + ".tar.bz2"
+        self.tar_gz_name = self.tar_basename + ".tar.gz"
+        self.tar_xz_name = self.tar_basename + ".tar.xz"
+        self.tar_zst_name = self.tar_basename + ".tar.zst"
         self.zip_name = self.tar_basename + ".zip"
         self.expected_files=['include/file1','LICENSES/test1.txt','autobuild-package.xml']
         self.expected_files.sort()
@@ -72,7 +75,10 @@ def tearDown(self):
         BaseTest.tearDown(self)
 
     def tar_has_expected(self,tar):
-        tarball = tarfile.open(tar, 'r:bz2')
+        if 'tar.zst' in tar:
+            tarball = common.ZstdTarFile(tar, 'r')
+        else:
+            tarball = tarfile.open(tar, 'r')
         packaged_files=tarball.getnames()
         packaged_files.sort()
         self.assertEqual(packaged_files, self.expected_files)
@@ -155,6 +161,27 @@ def test_autobuild_package(self):
         assert os.path.exists(self.tar_name), "%s does not exist" % self.tar_name
         self.tar_has_expected(self.tar_name)
         self.remove(self.tar_name)
+        self.autobuild("package",
+                       "--config-file=" + self.config_path,
+                       "--archive-format=tgz",
+                       "-p", "common")
+        assert os.path.exists(self.tar_gz_name), "%s does not exist" % self.tar_gz_name
+        self.tar_has_expected(self.tar_gz_name)
+        self.remove(self.tar_gz_name)
+        self.autobuild("package",
+                       "--config-file=" + self.config_path,
+                       "--archive-format=txz",
+                       "-p", "common")
+        assert os.path.exists(self.tar_xz_name), "%s does not exist" % self.tar_xz_name
+        self.tar_has_expected(self.tar_xz_name)
+        self.remove(self.tar_xz_name)
+        self.autobuild("package",
+                       "--config-file=" + self.config_path,
+                       "--archive-format=tzst",
+                       "-p", "common")
+        assert os.path.exists(self.tar_zst_name), "%s does not exist" % self.tar_zst_name
+        self.tar_has_expected(self.tar_zst_name)
+        self.remove(self.tar_zst_name)
         self.autobuild("package",
                        "--config-file=" + self.config_path,
                        "--archive-format=zip",