xsv: add --drop flag to partition command

The --drop flag drops the partition column itself.

PR #135 

2 files changed, 61 insertions, 2 deletions

diff --git a/src/cmd/partition.rs b/src/cmd/partition.rs
index c9fd7d4..68997a2 100644
--- a/src/cmd/partition.rs
+++ b/src/cmd/partition.rs
@@ -31,6 +31,7 @@ partition options:
     -p, --prefix-length <n>  Truncate the partition column after the
                            specified number of bytes when creating the
                            output file.
+    --drop                 Drop the partition column from results.
 
 Common options:
     -h, --help             Display this message
@@ -48,6 +49,7 @@ struct Args {
     arg_outdir: String,
     flag_filename: FilenameTemplate,
     flag_prefix_length: Option<usize>,
+    flag_drop: bool,
     flag_no_headers: bool,
     flag_delimiter: Option<Delimiter>,
 }
@@ -112,12 +114,22 @@ impl Args {
                     // We have a new key, so make a new writer.
                     let mut wtr = gen.writer(&*self.arg_outdir, key)?;
                     if !rconfig.no_headers {
-                        wtr.write_record(&headers)?;
+                        if self.flag_drop {
+                            wtr.write_record(headers.iter().enumerate()
+                                .filter_map(|(i, e)| if i != key_col { Some(e) } else { None }))?;
+                        } else {
+                            wtr.write_record(&headers)?;
+                        }
                     }
                     vacant.insert(wtr)
                 }
             };
-            wtr.write_byte_record(&row)?;
+            if self.flag_drop {
+                wtr.write_record(row.iter().enumerate()
+                    .filter_map(|(i, e)| if i != key_col { Some(e) } else { None }))?;
+            } else {
+                wtr.write_byte_record(&row)?;
+            }
         }
         Ok(())
     }
diff --git a/tests/test_partition.rs b/tests/test_partition.rs
index 5ed0cb5..4ff61e0 100644
--- a/tests/test_partition.rs
+++ b/tests/test_partition.rs
@@ -47,6 +47,31 @@ TX,Fort Worth
 }
 
 #[test]
+fn partition_drop() {
+    let wrk = Workdir::new("partition");
+    wrk.create("in.csv", data(true));
+
+    let mut cmd = wrk.command("partition");
+    cmd.arg("--drop").arg("state").arg(&wrk.path(".")).arg("in.csv");
+    wrk.run(&mut cmd);
+
+    part_eq!(wrk, "CA.csv", "\
+city
+San Francisco
+");
+    part_eq!(wrk, "NY.csv", "\
+city
+Manhatten
+Buffalo
+");
+    part_eq!(wrk, "TX.csv", "\
+city
+Dallas
+Fort Worth
+");
+}
+
+#[test]
 fn partition_without_headers() {
     let wrk = Workdir::new("partition_without_headers");
     wrk.create("in.csv", data(false));
@@ -69,6 +94,28 @@ TX,Fort Worth
 }
 
 #[test]
+fn partition_drop_without_headers() {
+    let wrk = Workdir::new("partition_without_headers");
+    wrk.create("in.csv", data(false));
+
+    let mut cmd = wrk.command("partition");
+    cmd.arg("--drop").arg("--no-headers").arg("1").arg(&wrk.path(".")).arg("in.csv");
+    wrk.run(&mut cmd);
+
+    part_eq!(wrk, "CA.csv", "\
+San Francisco
+");
+    part_eq!(wrk, "NY.csv", "\
+Manhatten
+Buffalo
+");
+    part_eq!(wrk, "TX.csv", "\
+Dallas
+Fort Worth
+");
+}
+
+#[test]
 fn partition_into_new_directory() {
     let wrk = Workdir::new("partition_into_new_directory");
     wrk.create("in.csv", data(true));