--keep option (#121)

src/headers.rs

@@ -1,3 +1,4 @@
+use std::collections::HashSet;
 use byteorder::{BigEndian, ReadBytesExt};
 use colors::{BitDepth, ColorType};
 use crc::crc32;
@@ -28,10 +29,12 @@ pub struct IhdrData {
 pub enum Headers {
     /// None
     None,
-    /// Some, with a list of 4-character chunk codes
+    /// Remove specific chunks
-    Some(Vec<String>),
+    Strip(Vec<String>),
     /// Headers that won't affect rendering (all but cHRM, gAMA, iCCP, sBIT, sRGB, bKGD, hIST, pHYs, sPLT)
     Safe,
+    /// Remove all non-critical chunks except these
+    Keep(HashSet<String>),
     /// All non-critical headers
     All,
 }

src/lib.rs

@@ -808,7 +808,12 @@ fn perform_strip(png: &mut PngData, opts: &Options) {
     match opts.strip {
         // Strip headers
         Headers::None => (),
-        Headers::Some(ref hdrs) => for hdr in hdrs {
+        Headers::Keep(ref hdrs) => {
+            png.aux_headers.retain(|chunk, _| {
+                hdrs.contains(chunk)
+            });
+        },
+        Headers::Strip(ref hdrs) => for hdr in hdrs {
             png.aux_headers.remove(hdr);
         },
         Headers::Safe => {

src/main.rs

@@ -94,6 +94,13 @@ fn main() {
             .takes_value(true)
             .value_name("mode")
             .conflicts_with("strip-safe"))
+        .arg(Arg::with_name("keep")
+            .help("Strip all optional metadata except objects in the comma-separated list")
+            .long("keep")
+            .takes_value(true)
+            .value_name("list")
+            .conflicts_with("strip")
+            .conflicts_with("strip-safe"))
         .arg(Arg::with_name("alpha")
             .help("Perform additional alpha optimizations")
             .short("a")
@@ -411,6 +418,12 @@ fn parse_opts_into_struct(matches: &ArgMatches) -> Result<(OutFile, Option<PathB
         opts.idat_recoding = false;
     }
 
+    if let Some(hdrs) = matches.value_of("keep") {
+        opts.strip = Headers::Keep(hdrs.split(',')
+            .map(|x| x.trim().to_owned())
+            .collect())
+    }
+
     if let Some(hdrs) = matches.value_of("strip") {
         let hdrs = hdrs.split(',')
             .map(|x| x.trim().to_owned())
@@ -433,7 +446,7 @@ fn parse_opts_into_struct(matches: &ArgMatches) -> Result<(OutFile, Option<PathB
                     return Err(format!("{} chunk is not allowed to be stripped", i));
                 }
             }
-            opts.strip = Headers::Some(hdrs);
+            opts.strip = Headers::Strip(hdrs);
         }
     }
 

tests/flags.rs

@@ -74,7 +74,7 @@ fn verbose_mode() {
 fn strip_headers_list() {
     let input = PathBuf::from("tests/files/strip_headers_list.png");
     let (output, mut opts) = get_opts(&input);
-    opts.strip = Headers::Some(vec!["iCCP".to_owned(), "tEXt".to_owned()]);
+    opts.strip = Headers::Strip(vec!["iCCP".to_owned(), "tEXt".to_owned()]);
 
     let png = PngData::new(&input, opts.fix_errors).unwrap();