Merge pull request #591 from pmeredit/topic/save_binary

Topic/save binary

src/commands/command.rs

@@ -496,6 +496,10 @@ pub trait WholeStreamCommand: Send + Sync {
         args: CommandArgs,
         registry: &registry::CommandRegistry,
     ) -> Result<OutputStream, ShellError>;
+
+    fn is_binary(&self) -> bool {
+        false
+    }
 }
 
 pub trait PerItemCommand: Send + Sync {
@@ -521,6 +525,10 @@ pub trait PerItemCommand: Send + Sync {
         raw_args: &RawCommandArgs,
         input: Tagged<Value>,
     ) -> Result<OutputStream, ShellError>;
+
+    fn is_binary(&self) -> bool {
+        false
+    }
 }
 
 pub enum Command {
@@ -608,6 +616,13 @@ impl Command {
             }
         }
     }
+
+    pub fn is_binary(&self) -> bool {
+        match self {
+            Command::WholeStream(command) => command.is_binary(),
+            Command::PerItem(command) => command.is_binary(),
+        }
+    }
 }
 
 pub struct FnFilterCommand {

src/commands/save.rs

@@ -6,6 +6,80 @@ use std::path::{Path, PathBuf};
 
 pub struct Save;
 
+macro_rules! process_string {
+    ($input:ident, $name_span:ident) => {{
+        let mut result_string = String::new();
+        for res in $input {
+            match res {
+                Tagged {
+                    item: Value::Primitive(Primitive::String(s)),
+                    ..
+                } => {
+                    result_string.push_str(&s);
+                }
+                _ => {
+                    yield core::task::Poll::Ready(Err(ShellError::labeled_error(
+                        "Save could not successfully save",
+                        "unexpected data during save",
+                        $name_span,
+                    )));
+                }
+            }
+        }
+        Ok(result_string.into_bytes())
+    }};
+}
+
+macro_rules! process_string_return_success {
+    ($result_vec:ident, $name_span:ident) => {{
+        let mut result_string = String::new();
+        for res in $result_vec {
+            match res {
+                Ok(ReturnSuccess::Value(Tagged {
+                    item: Value::Primitive(Primitive::String(s)),
+                    ..
+                })) => {
+                    result_string.push_str(&s);
+                }
+                _ => {
+                    yield core::task::Poll::Ready(Err(ShellError::labeled_error(
+                        "Save could not successfully save",
+                        "unexpected data during text save",
+                        $name_span,
+                    )));
+                }
+            }
+        }
+        Ok(result_string.into_bytes())
+    }};
+}
+
+macro_rules! process_binary_return_success {
+    ($result_vec:ident, $name_span:ident) => {{
+        let mut result_binary: Vec<u8> = Vec::new();
+        for res in $result_vec {
+            match res {
+                Ok(ReturnSuccess::Value(Tagged {
+                    item: Value::Binary(b),
+                    ..
+                })) => {
+                    for u in b.into_iter() {
+                        result_binary.push(u);
+                    }
+                }
+                _ => {
+                    yield core::task::Poll::Ready(Err(ShellError::labeled_error(
+                        "Save could not successfully save",
+                        "unexpected data during binary save",
+                        $name_span,
+                    )));
+                }
+            }
+        }
+        Ok(result_binary)
+    }};
+}
+
 #[derive(Deserialize)]
 pub struct SaveArgs {
     path: Option<Tagged<PathBuf>>,
@@ -96,7 +170,7 @@ fn save(
             }
         }
 
-        let content = if !save_raw {
+        let content : Result<Vec<u8>, ShellError> = if !save_raw {
             if let Some(extension) = full_path.extension() {
                 let command_name = format!("to-{}", extension.to_str().unwrap());
                 if let Some(converter) = registry.get_command(&command_name) {
@@ -116,60 +190,19 @@ fn save(
                     };
                     let mut result = converter.run(new_args.with_input(input), &registry);
                     let result_vec: Vec<Result<ReturnSuccess, ShellError>> = result.drain_vec().await;
-                    let mut result_string = String::new();
-                    for res in result_vec {
-                        match res {
-                            Ok(ReturnSuccess::Value(Tagged { item: Value::Primitive(Primitive::String(s)), .. })) => {
-                                result_string.push_str(&s);
-                            }
-                            _ => {
-                                yield Err(ShellError::labeled_error(
-                                    "Save could not successfully save",
-                                    "unexpected data during save",
-                                    name_span,
-                                ));
-                            },
-                        }
+                    if converter.is_binary() {
+                        process_binary_return_success!(result_vec, name_span)
+                    } else {
+                        process_string_return_success!(result_vec, name_span)
                     }
-                    Ok(result_string)
                 } else {
-                    let mut result_string = String::new();
-                    for res in input {
-                        match res {
-                            Tagged { item: Value::Primitive(Primitive::String(s)), .. } => {
-                                result_string.push_str(&s);
-                            }
-                            _ => {
-                                yield Err(ShellError::labeled_error(
-                                    "Save could not successfully save",
-                                    "unexpected data during save",
-                                    name_span,
-                                ));
-                            },
-                        }
-                    }
-                    Ok(result_string)
+                    process_string!(input, name_span)
                 }
             } else {
-                let mut result_string = String::new();
-                for res in input {
-                    match res {
-                        Tagged { item: Value::Primitive(Primitive::String(s)), .. } => {
-                            result_string.push_str(&s);
-                        }
-                        _ => {
-                            yield Err(ShellError::labeled_error(
-                                "Save could not successfully save",
-                                "unexpected data during save",
-                                name_span,
-                            ));
-                        },
-                    }
-                }
-                Ok(result_string)
+                process_string!(input, name_span)
             }
         } else {
-            string_from(&input)
+            Ok(string_from(&input).into_bytes())
         };
 
         match content {
@@ -185,7 +218,7 @@ fn save(
     Ok(OutputStream::new(stream))
 }
 
-fn string_from(input: &Vec<Tagged<Value>>) -> Result<String, ShellError> {
+fn string_from(input: &Vec<Tagged<Value>>) -> String {
     let mut save_data = String::new();
 
     if input.len() > 0 {
@@ -202,5 +235,5 @@ fn string_from(input: &Vec<Tagged<Value>>) -> Result<String, ShellError> {
         }
     }
 
-    Ok(save_data)
+    save_data
 }

src/commands/to_bson.rs

@@ -26,6 +26,10 @@ impl WholeStreamCommand for ToBSON {
     ) -> Result<OutputStream, ShellError> {
         to_bson(args, registry)
     }
+
+    fn is_binary(&self) -> bool {
+        true
+    }
 }
 
 pub fn value_to_bson_value(v: &Tagged<Value>) -> Result<Bson, ShellError> {

src/commands/to_sqlite.rs

@@ -27,6 +27,10 @@ impl WholeStreamCommand for ToSQLite {
     ) -> Result<OutputStream, ShellError> {
         to_sqlite(args, registry)
     }
+
+    fn is_binary(&self) -> bool {
+        true
+    }
 }
 
 pub struct ToDB;
@@ -51,6 +55,10 @@ impl WholeStreamCommand for ToDB {
     ) -> Result<OutputStream, ShellError> {
         to_sqlite(args, registry)
     }
+
+    fn is_binary(&self) -> bool {
+        true
+    }
 }
 
 fn comma_concat(acc: String, current: String) -> String {

tests/commands_test.rs

@@ -36,8 +36,8 @@ fn save_figures_out_intelligently_where_to_write_out_with_metadata() {
                 description = "A shell for the GitHub era"
                 license = "ISC"
                 edition = "2018"
-            "#)
-        ]);
+            "#,
+        )]);
 
         let subject_file = dirs.test().join("cargo_sample.toml");
 
@@ -66,3 +66,32 @@ fn save_can_write_out_csv() {
         assert!(actual.contains("[list list],A shell for the GitHub era,2018,ISC,nu,0.2.0"));
     })
 }
+
+#[test]
+fn save_can_write_out_bson() {
+    Playground::setup("save_test_3", |dirs, _| {
+        let expected_file = dirs.test().join("cargo_sample.bson");
+
+        nu!(
+            cwd: dirs.root(),
+            "open {}/cargo_sample.toml | inc package.version --minor | get package | save save_test_3/cargo_sample.bson",
+            dirs.formats()
+        );
+
+        let actual = h::file_contents_binary(expected_file);
+        assert!(
+            actual
+                == vec![
+                    168, 0, 0, 0, 4, 97, 117, 116, 104, 111, 114, 115, 0, 43, 0, 0, 0, 2, 48, 0,
+                    31, 0, 0, 0, 89, 101, 104, 117, 100, 97, 32, 75, 97, 116, 122, 32, 60, 119,
+                    121, 99, 97, 116, 115, 64, 103, 109, 97, 105, 108, 46, 99, 111, 109, 62, 0, 0,
+                    2, 100, 101, 115, 99, 114, 105, 112, 116, 105, 111, 110, 0, 27, 0, 0, 0, 65,
+                    32, 115, 104, 101, 108, 108, 32, 102, 111, 114, 32, 116, 104, 101, 32, 71, 105,
+                    116, 72, 117, 98, 32, 101, 114, 97, 0, 2, 101, 100, 105, 116, 105, 111, 110, 0,
+                    5, 0, 0, 0, 50, 48, 49, 56, 0, 2, 108, 105, 99, 101, 110, 115, 101, 0, 4, 0, 0,
+                    0, 73, 83, 67, 0, 2, 110, 97, 109, 101, 0, 3, 0, 0, 0, 110, 117, 0, 2, 118,
+                    101, 114, 115, 105, 111, 110, 0, 6, 0, 0, 0, 48, 46, 50, 46, 48, 0, 0
+                ]
+        );
+    })
+}

tests/helpers/mod.rs

@@ -307,6 +307,14 @@ pub fn file_contents(full_path: impl AsRef<Path>) -> String {
     contents
 }
 
+pub fn file_contents_binary(full_path: impl AsRef<Path>) -> Vec<u8> {
+    let mut file = std::fs::File::open(full_path.as_ref()).expect("can not open file");
+    let mut contents = Vec::new();
+    file.read_to_end(&mut contents)
+        .expect("can not read file");
+    contents
+}
+
 pub fn line_ending() -> String {
     #[cfg(windows)]
     {