use nu_engine::CallExt; use nu_protocol::ast::Call; use nu_protocol::engine::{Command, EngineState, Stack}; use nu_protocol::{ Category, Example, PipelineData, RawStream, ShellError, Signature, Span, Spanned, SyntaxShape, Type, Value, }; use std::fs::File; use std::io::Write; use std::path::Path; use std::thread; use crate::progress_bar; #[derive(Clone)] pub struct Save; impl Command for Save { fn name(&self) -> &str { "save" } fn usage(&self) -> &str { "Save a file." } fn search_terms(&self) -> Vec<&str> { vec![ "write", "write_file", "append", "redirection", "file", "io", ">", ">>", ] } fn signature(&self) -> nu_protocol::Signature { Signature::build("save") .input_output_types(vec![(Type::Any, Type::Nothing)]) .required("filename", SyntaxShape::Filepath, "the filename to use") .named( "stderr", SyntaxShape::Filepath, "the filename used to save stderr, only works with `-r` flag", Some('e'), ) .switch("raw", "save file as raw binary", Some('r')) .switch("append", "append input to the end of the file", Some('a')) .switch("force", "overwrite the destination", Some('f')) .switch("progress", "enable progress bar", Some('p')) .category(Category::FileSystem) } fn run( &self, engine_state: &EngineState, stack: &mut Stack, call: &Call, input: PipelineData, ) -> Result { let raw = call.has_flag("raw"); let append = call.has_flag("append"); let force = call.has_flag("force"); let progress = call.has_flag("progress"); let span = call.head; let path = call.req::>(engine_state, stack, 0)?; let stderr_path = call.get_flag::>(engine_state, stack, "stderr")?; match input { PipelineData::ExternalStream { stdout: None, .. } => { // Open files to possibly truncate them let _ = get_files(&path, &stderr_path, append, force)?; Ok(PipelineData::empty()) } PipelineData::ExternalStream { stdout: Some(stream), stderr, .. } => { let (file, stderr_file) = get_files(&path, &stderr_path, append, force)?; // delegate a thread to redirect stderr to result. let handler = stderr.map(|stderr_stream| match stderr_file { Some(stderr_file) => thread::Builder::new() .name("stderr redirector".to_string()) .spawn(move || stream_to_file(stderr_stream, stderr_file, span, progress)) .expect("Failed to create thread"), None => thread::Builder::new() .name("stderr redirector".to_string()) .spawn(move || { let _ = stderr_stream.into_bytes(); Ok(PipelineData::empty()) }) .expect("Failed to create thread"), }); let res = stream_to_file(stream, file, span, progress); if let Some(h) = handler { h.join().map_err(|err| ShellError::ExternalCommand { label: "Fail to receive external commands stderr message".to_string(), help: format!("{err:?}"), span, })??; res } else { res } } PipelineData::ListStream(ls, _) if raw || prepare_path(&path, append, force)?.0.extension().is_none() => { let (mut file, _) = get_files(&path, &stderr_path, append, force)?; for val in ls { file.write_all(&value_to_bytes(val)?) .map_err(|err| ShellError::IOError(err.to_string()))?; file.write_all("\n".as_bytes()) .map_err(|err| ShellError::IOError(err.to_string()))?; } file.flush()?; Ok(PipelineData::empty()) } input => { let bytes = input_to_bytes(input, Path::new(&path.item), raw, engine_state, stack, span)?; // Only open file after successful conversion let (mut file, _) = get_files(&path, &stderr_path, append, force)?; file.write_all(&bytes) .map_err(|err| ShellError::IOError(err.to_string()))?; file.flush()?; Ok(PipelineData::empty()) } } } fn examples(&self) -> Vec { vec![ Example { description: "Save a string to foo.txt in the current directory", example: r#"'save me' | save foo.txt"#, result: None, }, Example { description: "Append a string to the end of foo.txt", example: r#"'append me' | save --append foo.txt"#, result: None, }, Example { description: "Save a record to foo.json in the current directory", example: r#"{ a: 1, b: 2 } | save foo.json"#, result: None, }, Example { description: "Save a running program's stderr to foo.txt", example: r#"do -i {} | save foo.txt --stderr foo.txt"#, result: None, }, Example { description: "Save a running program's stderr to separate file", example: r#"do -i {} | save foo.txt --stderr bar.txt"#, result: None, }, ] } } /// Convert [`PipelineData`] bytes to write in file, possibly converting /// to format of output file fn input_to_bytes( input: PipelineData, path: &Path, raw: bool, engine_state: &EngineState, stack: &mut Stack, span: Span, ) -> Result, ShellError> { let ext = if raw { None // if is extern stream , in other words , not value } else if let PipelineData::ExternalStream { .. } = input { None } else if let PipelineData::Value(Value::String { .. }, ..) = input { None } else { path.extension() .map(|name| name.to_string_lossy().to_string()) }; if let Some(ext) = ext { convert_to_extension(engine_state, &ext, stack, input, span) } else { let value = input.into_value(span); value_to_bytes(value) } } /// Convert given data into content of file of specified extension if /// corresponding `to` command exists. Otherwise attempt to convert /// data to bytes as is fn convert_to_extension( engine_state: &EngineState, extension: &str, stack: &mut Stack, input: PipelineData, span: Span, ) -> Result, ShellError> { let converter = engine_state.find_decl(format!("to {extension}").as_bytes(), &[]); let output = match converter { Some(converter_id) => { let output = engine_state.get_decl(converter_id).run( engine_state, stack, &Call::new(span), input, )?; output.into_value(span) } None => input.into_value(span), }; value_to_bytes(output) } /// Convert [`Value::String`] [`Value::Binary`] or [`Value::List`] into [`Vec`] of bytes /// /// Propagates [`Value::Error`] and creates error otherwise fn value_to_bytes(value: Value) -> Result, ShellError> { match value { Value::String { val, .. } => Ok(val.into_bytes()), Value::Binary { val, .. } => Ok(val), Value::List { vals, .. } => { let val = vals .into_iter() .map(|it| it.as_string()) .collect::, ShellError>>()? .join("\n") + "\n"; Ok(val.into_bytes()) } // Propagate errors by explicitly matching them before the final case. Value::Error { error } => Err(*error), other => Ok(other.as_string()?.into_bytes()), } } /// Convert string path to [`Path`] and [`Span`] and check if this path /// can be used with given flags fn prepare_path( path: &Spanned, append: bool, force: bool, ) -> Result<(&Path, Span), ShellError> { let span = path.span; let path = Path::new(&path.item); if !(force || append) && path.exists() { Err(ShellError::GenericError( "Destination file already exists".into(), format!( "Destination file '{}' already exists", path.to_string_lossy() ), Some(span), Some("you can use -f, --force to force overwriting the destination".into()), Vec::new(), )) } else { Ok((path, span)) } } fn open_file(path: &Path, span: Span, append: bool) -> Result { let file = match (append, path.exists()) { (true, true) => std::fs::OpenOptions::new() .write(true) .append(true) .open(path), _ => std::fs::File::create(path), }; file.map_err(|err| { ShellError::GenericError( "Permission denied".into(), err.to_string(), Some(span), None, Vec::new(), ) }) } /// Get output file and optional stderr file fn get_files( path: &Spanned, stderr_path: &Option>, append: bool, force: bool, ) -> Result<(File, Option), ShellError> { // First check both paths let (path, path_span) = prepare_path(path, append, force)?; let stderr_path_and_span = stderr_path .as_ref() .map(|stderr_path| prepare_path(stderr_path, append, force)) .transpose()?; // Only if both files can be used open and possibly truncate them let file = open_file(path, path_span, append)?; let stderr_file = stderr_path_and_span .map(|(stderr_path, stderr_path_span)| { if path == stderr_path { Err(ShellError::GenericError( "input and stderr input to same file".to_string(), "can't save both input and stderr input to the same file".to_string(), Some(stderr_path_span), Some("you should use `o+e> file` instead".to_string()), vec![], )) } else { open_file(stderr_path, stderr_path_span, append) } }) .transpose()?; Ok((file, stderr_file)) } fn stream_to_file( mut stream: RawStream, file: File, span: Span, progress: bool, ) -> Result { // https://github.com/nushell/nushell/pull/9377 contains the reason // for not using BufWriter let mut writer = file; let mut bytes_processed: u64 = 0; let bytes_processed_p = &mut bytes_processed; let file_total_size = stream.known_size; let mut process_failed = false; let process_failed_p = &mut process_failed; // Create the progress bar // It looks a bit messy but I am doing it this way to avoid // creating the bar when is not needed let (mut bar_opt, bar_opt_clone) = if progress { let tmp_bar = progress_bar::NuProgressBar::new(file_total_size); let tmp_bar_clone = tmp_bar.clone(); (Some(tmp_bar), Some(tmp_bar_clone)) } else { (None, None) }; let result = stream .try_for_each(move |result| { let buf = match result { Ok(v) => match v { Value::String { val, .. } => val.into_bytes(), Value::Binary { val, .. } => val, // Propagate errors by explicitly matching them before the final case. Value::Error { error } => return Err(*error), other => { return Err(ShellError::OnlySupportsThisInputType { exp_input_type: "string or binary".into(), wrong_type: other.get_type().to_string(), dst_span: span, src_span: other.expect_span(), }); } }, Err(err) => { *process_failed_p = true; return Err(err); } }; // If the `progress` flag is set then if progress { // Update the total amount of bytes that has been saved and then print the progress bar *bytes_processed_p += buf.len() as u64; if let Some(bar) = &mut bar_opt { bar.update_bar(*bytes_processed_p); } } if let Err(err) = writer.write(&buf) { *process_failed_p = true; return Err(ShellError::IOError(err.to_string())); } Ok(()) }) .map(|_| PipelineData::empty()); // If the `progress` flag is set then if progress { // If the process failed, stop the progress bar with an error message. if process_failed { if let Some(bar) = bar_opt_clone { bar.abandoned_msg("# Error while saving #".to_owned()); } } } // And finally return the stream result. result }