Add pigz support

src/icepack/__init__.py

@@ -9,7 +9,7 @@
 from pydantic import ValidationError
 
 from icepack.error import InvalidArchiveError
-from icepack.helper import Age, File, SSH, Zip
+from icepack.helper import Age, File, GZip, SSH, Zip
 from icepack.meta import SECRET_KEY, PUBLIC_KEY, ALLOWED_SIGNERS
 from icepack.model import Checksum, Compression, Encryption
 from icepack.model import DirEntry, FileEntry, Metadata
@@ -146,16 +146,15 @@ def _load_metadata(self):
 
     def _uncompress_path(self, src_path, dst_path, compression):
         """Uncompressed src_path to dst_path."""
-        with open(dst_path, 'wb') as dst:
-            with open(src_path, 'rb') as src_file:
-                if compression == Compression.BZ2:
+        if compression == Compression.BZ2:
+            with open(dst_path, 'wb') as dst:
+                with open(src_path, 'rb') as src_file:
                     with bz2.open(src_file, 'rb') as src:
                         copyfileobj(src, dst, _BUFFER_SIZE)
-                elif compression == Compression.GZ:
-                    with gzip.open(src_file, 'rb') as src:
-                        copyfileobj(src, dst, _BUFFER_SIZE)
-                else:
-                    raise Exception('Unsupported compression.')
+        elif compression == Compression.GZ:
+            GZip.decompress(src_path, dst_path)
+        else:
+            raise Exception('Unsupported compression.')
 
 
 class IcepackWriter(IcepackBase):
@@ -249,15 +248,14 @@ def add_metadata(self):
     def _compress_path(self, src_path):
         """Return the temporary Path of the compressed src_path."""
         tmp_path = self._mktemp()
-        with open(src_path, 'rb') as src:
-            if self._compression == Compression.BZ2:
+        if self._compression == Compression.BZ2:
+            with open(src_path, 'rb') as src:
                 with bz2.open(tmp_path, 'wb') as dst:
                     copyfileobj(src, dst, _BUFFER_SIZE)
-            elif self._compression == Compression.GZ:
-                with gzip.open(tmp_path, 'wb') as dst:
-                    copyfileobj(src, dst, _BUFFER_SIZE)
-            else:
-                raise Exception('Unsupported compression.')
+        elif self._compression == Compression.GZ:
+            GZip.compress(src_path, tmp_path)
+        else:
+            raise Exception('Unsupported compression.')
         return tmp_path
 
     def _encrypt_path(self, src_path):

src/icepack/cli.py

@@ -3,7 +3,7 @@
 import click
 
 from icepack import IcepackReader, create_archive, extract_archive
-from icepack.helper import Age, File, SSH
+from icepack.helper import Age, File, GZip, SSH
 from icepack.meta import NAME, VERSION, SECRET_KEY, PUBLIC_KEY
 from icepack.model import Compression
 
@@ -167,6 +167,9 @@ def version(ctx, dependencies):
         click.echo(f'✅ ssh-keygen found.')
     else:
         click.echo(f'❌ ssh-keygen not found.')
+    pigz_version = GZip.pigz_version()
+    if pigz_version:
+        click.echo(f'✅ pigz found. (Version: {pigz_version})')
 
 
 @icepack.group()

src/icepack/helper.py

@@ -1,3 +1,4 @@
+import gzip
 import hashlib
 import os
 from pathlib import Path
@@ -11,6 +12,9 @@
 from icepack.meta import NAME, SECRET_KEY, PUBLIC_KEY, ALLOWED_SIGNERS
 
 
+# If this environment variable is set to 'false', do not use pigz.
+PIGZ_ENV = 'ICEPACK_PIGZ'
+
 _BUFFER_SIZE = 64 * 1024
 _PUBLIC_KEY_PREFIX = 'age'
 _SECRET_KEY_PREFIX = 'AGE-SECRET-KEY-'  # nosec No secret
@@ -136,6 +140,50 @@ def sha256(path):
         return d.hexdigest()
 
 
+class GZip():
+    """gzip helpers."""
+    @staticmethod
+    def compress(src_path, dst_path):
+        """Compress src_path to dst_path."""
+        pigz_env = os.environ.get(PIGZ_ENV)
+        if GZip.has_pigz() and pigz_env != 'false':
+            with open(dst_path, 'wb') as dst:
+                subprocess.run(  # nosec Trusted input
+                    ['pigz', '-c', '-m', '-n', '-9', str(src_path)],
+                    stdout=dst,
+                    check=True)
+        else:
+            with open(src_path, 'rb') as src:
+                with gzip.open(dst_path, 'wb') as dst:
+                    copyfileobj(src, dst, _BUFFER_SIZE)
+
+    @staticmethod
+    def decompress(src_path, dst_path):
+        """Decompress src_path to dst_path."""
+        with gzip.open(src_path, 'rb') as src:
+            with open(dst_path, 'wb') as dst:
+                copyfileobj(src, dst, _BUFFER_SIZE)
+
+    @staticmethod
+    def has_pigz():
+        """Return True if pigz is available."""
+        return which('pigz') is not None
+
+    @staticmethod
+    def pigz_version():
+        """Return the pigz version, or None."""
+        version = None
+        if GZip.has_pigz():
+            result = subprocess.run(  # nosec Trusted input
+                ['pigz', '--version'],
+                capture_output=True,
+                text=True,
+                timeout=5)
+            if result.returncode == 0:
+                version = result.stdout.strip()
+        return version
+
+
 class SSH():
     """ssh-keygen helpers."""
 

tests/test_helper_gzip.py

@@ -0,0 +1,40 @@
+import os
+
+import pytest
+
+from icepack.helper import PIGZ_ENV, File, GZip
+
+from helper import src_path, dst_path
+
+
+@pytest.fixture
+def pigz_env():
+    """Preserve the existing PIGZ_ENV value, if any."""
+    old_value = os.environ.get(PIGZ_ENV)
+    yield
+    if old_value is not None:
+        os.environ[PIGZ_ENV] = old_value
+    elif PIGZ_ENV in os.environ:
+        del os.environ[PIGZ_ENV]
+
+
+def test_ensure_pigz():
+    """Ensure that pigz is available."""
+    assert GZip.has_pigz() is True
+    assert GZip.pigz_version is not None
+
+
+def test_without_pigz(src_path, dst_path, pigz_env):
+    """Test compression without pigz."""
+    os.environ[PIGZ_ENV] = 'false'
+    GZip.compress(src_path / 'foo', dst_path / 'foo.gz')
+    GZip.decompress(dst_path / 'foo.gz', dst_path / 'foo')
+    assert File.sha256(dst_path / 'foo') == 'b5bb9d8014a0f9b1d61e21e796d78dccdf1352f23cd32812f4850b878ae4944c'  # noqa
+
+
+def test_with_pigz(src_path, dst_path, pigz_env):
+    """Test compression with pigz."""
+    os.environ[PIGZ_ENV] = 'true'
+    GZip.compress(src_path / 'foo', dst_path / 'foo.gz')
+    GZip.decompress(dst_path / 'foo.gz', dst_path / 'foo')
+    assert File.sha256(dst_path / 'foo') == 'b5bb9d8014a0f9b1d61e21e796d78dccdf1352f23cd32812f4850b878ae4944c'  # noqa