From 15514f9cd76c8f4dc7ddadbe0d62f7285618b32d Mon Sep 17 00:00:00 2001
From: Dusty Mabe <dusty@dustymabe.com>
Date: Wed, 31 Jan 2024 15:32:52 -0500
Subject: [PATCH] cmd-compress: support zstd as an option

I implemented this to investigate it as an option for
https://github.com/coreos/fedora-coreos-tracker/issues/1660
so figured I may as well post the code up for inclusion.
---
 src/cmd-compress | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/src/cmd-compress b/src/cmd-compress
index 1e01d1263c..4fb6ea8c13 100755
--- a/src/cmd-compress
+++ b/src/cmd-compress
@@ -32,7 +32,7 @@ parser = argparse.ArgumentParser()
 parser.add_argument("--build", help="Build ID")
 parser.add_argument("--artifact", default=[], action='append',
                     help="Only compress given image ARTIFACT", metavar="ARTIFACT")
-parser.add_argument("--compressor", choices=['xz', 'gzip'],
+parser.add_argument("--compressor", choices=['xz', 'gzip', 'zstd'],
                     help="Override compressor to use")
 parser.add_argument("--mode",
                     choices=['compress', 'uncompress'],
@@ -53,7 +53,7 @@ else:
 print(f"Targeting build: {build}")
 
 # common extensions for known compressors
-ext_dict = {'xz': '.xz', 'gzip': '.gz'}
+ext_dict = {'xz': '.xz', 'gzip': '.gz', 'zstd': '.zst'}
 
 
 def get_cpu_param(param):
@@ -115,9 +115,11 @@ def compress_one_builddir(builddir):
             img['uncompressed-sha256'] = img['sha256']
             img['uncompressed-size'] = img['size']
             with open(tmpfile, 'wb') as f:
+                t = ncpu()
                 if args.compressor == 'xz':
-                    t = ncpu()
                     runcmd(['xz', '-c9', f'-T{t}', filepath], stdout=f)
+                elif args.compressor == 'zstd':
+                    runcmd(['zstd', '-19', '-c', f'-T{t}', filepath], stdout=f)
                 else:
                     runcmd(['gzip', f'-{gzip_level}', '-c', filepath], stdout=f)
             file_with_ext = file + ext
@@ -196,6 +198,8 @@ def uncompress_one_builddir(builddir):
                 if file.endswith('xz'):
                     t = ncpu()
                     runcmd(['xz', '-dc', f'-T{t}', filepath], stdout=f)
+                elif file.endswith('zst'):
+                    runcmd(['zstd', '-dc', filepath], stdout=f)
                 elif file.endswith('gz'):
                     runcmd(['gzip', '-dc', filepath], stdout=f)
                 else: