generate: allow choosing the compression formats to output

6 files changed, 99 insertions, 8 deletions

diff --git a/src/compression.rs b/src/compression.rs
index b8cfd9a..b3010cb 100644
--- a/src/compression.rs
+++ b/src/compression.rs
@@ -1,10 +1,11 @@
 use anyhow::{Context, Error};
 use flate2::{read::GzDecoder, write::GzEncoder};
 use rayon::prelude::*;
-use std::{io::Read, io::Write, path::Path};
+use std::{convert::TryFrom, io::Read, io::Write, path::Path};
 use xz2::{read::XzDecoder, write::XzEncoder};
 
-pub(crate) enum CompressionFormat {
+#[derive(Debug, Copy, Clone)]
+pub enum CompressionFormat {
     Gz,
     Xz,
 }
@@ -59,6 +60,38 @@ impl CompressionFormat {
     }
 }
 
+/// This struct wraps Vec<CompressionFormat> in order to parse the value from the command line.
+#[derive(Debug, Clone)]
+pub struct CompressionFormats(Vec<CompressionFormat>);
+
+impl TryFrom<&'_ str> for CompressionFormats {
+    type Error = Error;
+
+    fn try_from(value: &str) -> Result<Self, Self::Error> {
+        let mut parsed = Vec::new();
+        for format in value.split(',') {
+            match format.trim() {
+                "gz" => parsed.push(CompressionFormat::Gz),
+                "xz" => parsed.push(CompressionFormat::Xz),
+                other => anyhow::bail!("unknown compression format: {}", other),
+            }
+        }
+        Ok(CompressionFormats(parsed))
+    }
+}
+
+impl Default for CompressionFormats {
+    fn default() -> Self {
+        Self(vec![CompressionFormat::Gz, CompressionFormat::Xz])
+    }
+}
+
+impl CompressionFormats {
+    pub(crate) fn iter(&self) -> impl Iterator<Item = CompressionFormat> + '_ {
+        self.0.iter().map(|i| *i)
+    }
+}
+
 pub(crate) trait Encoder: Send + Write {
     fn finish(self: Box<Self>) -> Result<(), Error>;
 }
diff --git a/src/generator.rs b/src/generator.rs
index 99ba5f1..2601eb5 100644
--- a/src/generator.rs
+++ b/src/generator.rs
@@ -1,5 +1,6 @@
 use super::Scripter;
 use super::Tarballer;
+use crate::compression::CompressionFormats;
 use crate::util::*;
 use anyhow::{bail, format_err, Context, Result};
 use std::io::Write;
@@ -40,6 +41,9 @@ actor! {
 
         /// The location to put the final image and tarball
         output_dir: String = "./dist",
+
+        /// The formats used to compress the tarball
+        compression_formats: CompressionFormats = CompressionFormats::default(),
     }
 }
 
@@ -95,7 +99,8 @@ impl Generator {
         tarballer
             .work_dir(self.work_dir)
             .input(self.package_name)
-            .output(path_to_str(&output)?.into());
+            .output(path_to_str(&output)?.into())
+            .compression_formats(self.compression_formats.clone());
         tarballer.run()?;
 
         Ok(())
diff --git a/src/main.rs b/src/main.rs
index 9615c9d..c5d9f23 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -60,6 +60,7 @@ fn generate(matches: &ArgMatches<'_>) -> Result<()> {
         "image-dir" => image_dir,
         "work-dir" => work_dir,
         "output-dir" => output_dir,
+        "compression-formats" => compression_formats,
     });
 
     generator.run().context("failed to generate installer")?;
diff --git a/src/main.yml b/src/main.yml
index 2f9978b..afe235c 100644
--- a/src/main.yml
+++ b/src/main.yml
@@ -60,6 +60,11 @@ subcommands:
             long: output-dir
             takes_value: true
             value_name: DIR
+        - compression-formats:
+            help: Comma-separated list of compression formats to use
+            long: compression-formats
+            takes_value: true
+            value_name: FORMAT
   - combine:
       about: Combine installer tarballs
       args:
diff --git a/src/tarballer.rs b/src/tarballer.rs
index c23cac1..4ac8cf7 100644
--- a/src/tarballer.rs
+++ b/src/tarballer.rs
@@ -5,7 +5,10 @@ use std::path::Path;
 use tar::{Builder, Header};
 use walkdir::WalkDir;
 
-use crate::{compression::CombinedEncoder, compression::CompressionFormat, util::*};
+use crate::{
+    compression::{CombinedEncoder, CompressionFormats},
+    util::*,
+};
 
 actor! {
     #[derive(Debug)]
@@ -18,6 +21,9 @@ actor! {
 
         /// The folder in which the input is to be found.
         work_dir: String = "./workdir",
+
+        /// The formats used to compress the tarball.
+        compression_formats: CompressionFormats = CompressionFormats::default(),
     }
 }
 
@@ -25,10 +31,12 @@ impl Tarballer {
     /// Generates the actual tarballs
     pub fn run(self) -> Result<()> {
         let tarball_name = self.output.clone() + ".tar";
-        let encoder = CombinedEncoder::new(vec![
-            CompressionFormat::Gz.encode(&tarball_name)?,
-            CompressionFormat::Xz.encode(&tarball_name)?,
-        ]);
+        let encoder = CombinedEncoder::new(
+            self.compression_formats
+                .iter()
+                .map(|f| f.encode(&tarball_name))
+                .collect::<Result<Vec<_>>>()?,
+        );
 
         // Sort files by their suffix, to group files with the same name from
         // different locations (likely identical) and files with the same
diff --git a/test.sh b/test.sh
index 3dc5e49..a73ea9d 100755
--- a/test.sh
+++ b/test.sh
@@ -1190,6 +1190,45 @@ combine_installers_different_input_compression_formats() {
 }
 runtest combine_installers_different_input_compression_formats
 
+generate_compression_formats_one() {
+    try sh "$S/gen-installer.sh" \
+        --image-dir="$TEST_DIR/image1" \
+        --work-dir="$WORK_DIR" \
+        --output-dir="$OUT_DIR" \
+        --package-name="rustc" \
+        --component-name="rustc" \
+        --compression-formats="xz"
+
+    try test ! -e "${OUT_DIR}/rustc.tar.gz"
+    try test -e "${OUT_DIR}/rustc.tar.xz"
+}
+runtest generate_compression_formats_one
+
+generate_compression_formats_multiple() {
+    try sh "$S/gen-installer.sh" \
+        --image-dir="$TEST_DIR/image1" \
+        --work-dir="$WORK_DIR" \
+        --output-dir="$OUT_DIR" \
+        --package-name="rustc" \
+        --component-name="rustc" \
+        --compression-formats="gz,xz"
+
+    try test -e "${OUT_DIR}/rustc.tar.gz"
+    try test -e "${OUT_DIR}/rustc.tar.xz"
+}
+runtest generate_compression_formats_multiple
+
+generate_compression_formats_error() {
+    expect_fail sh "$S/gen-installer.sh" \
+        --image-dir="$TEST_DIR/image1" \
+        --work-dir="$WORK_DIR" \
+        --output-dir="$OUT_DIR" \
+        --package-name="rustc" \
+        --component-name="rustc" \
+        --compression-formats="xz,foobar"
+}
+runtest generate_compression_formats_error
+
 echo
 echo "TOTAL SUCCESS!"
 echo