use crate::progress_bar; #[allow(deprecated)] use nu_engine::{command_prelude::*, current_dir}; use nu_path::expand_path_with; use nu_protocol::{ ast::{Expr, Expression}, DataSource, OutDest, PipelineMetadata, RawStream, }; use std::{ fs::File, io::Write, path::{Path, PathBuf}, thread, }; #[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(engine_state, stack, "raw")?; let append = call.has_flag(engine_state, stack, "append")?; let force = call.has_flag(engine_state, stack, "force")?; let progress = call.has_flag(engine_state, stack, "progress")?; let out_append = if let Some(Expression { expr: Expr::Bool(out_append), .. }) = call.get_parser_info("out-append") { *out_append } else { false }; let err_append = if let Some(Expression { expr: Expr::Bool(err_append), .. }) = call.get_parser_info("err-append") { *err_append } else { false }; let span = call.head; #[allow(deprecated)] let cwd = current_dir(engine_state, stack)?; let path_arg = call.req::>(engine_state, stack, 0)?; let path = Spanned { item: expand_path_with(path_arg.item, &cwd, true), span: path_arg.span, }; let stderr_path = call .get_flag::>(engine_state, stack, "stderr")? .map(|arg| Spanned { item: expand_path_with(arg.item, cwd, true), span: arg.span, }); match input { PipelineData::ExternalStream { stdout, stderr, metadata, .. } => { check_saving_to_source_file(metadata.as_ref(), &path, stderr_path.as_ref())?; let (file, stderr_file) = get_files( &path, stderr_path.as_ref(), append, out_append, err_append, force, )?; match (stdout, stderr) { (Some(stdout), stderr) => { // delegate a thread to redirect stderr to result. let handler = stderr .map(|stderr| match stderr_file { Some(stderr_file) => thread::Builder::new() .name("stderr redirector".to_string()) .spawn(move || { stream_to_file(stderr, stderr_file, span, progress) }), None => thread::Builder::new() .name("stderr redirector".to_string()) .spawn(move || stderr.drain()), }) .transpose() .err_span(span)?; let res = stream_to_file(stdout, 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?; } (None, Some(stderr)) => match stderr_file { Some(stderr_file) => stream_to_file(stderr, stderr_file, span, progress)?, None => stderr.drain()?, }, (None, None) => {} }; Ok(PipelineData::Empty) } PipelineData::ListStream(ls, pipeline_metadata) if raw || prepare_path(&path, append, force)?.0.extension().is_none() => { check_saving_to_source_file( pipeline_metadata.as_ref(), &path, stderr_path.as_ref(), )?; let (mut file, _) = get_files( &path, stderr_path.as_ref(), append, out_append, err_append, force, )?; for val in ls { file.write_all(&value_to_bytes(val)?) .map_err(|err| ShellError::IOError { msg: err.to_string(), })?; file.write_all("\n".as_bytes()) .map_err(|err| ShellError::IOError { msg: err.to_string(), })?; } file.flush()?; Ok(PipelineData::empty()) } input => { check_saving_to_source_file( input.metadata().as_ref(), &path, stderr_path.as_ref(), )?; 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.as_ref(), append, out_append, err_append, force, )?; file.write_all(&bytes).map_err(|err| ShellError::IOError { msg: 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, }, ] } fn pipe_redirection(&self) -> (Option, Option) { (Some(OutDest::Capture), Some(OutDest::Capture)) } } fn saving_to_source_file_error(dest: &Spanned) -> ShellError { ShellError::GenericError { error: "pipeline input and output are the same file".into(), msg: format!( "can't save output to '{}' while it's being read", dest.item.display() ), span: Some(dest.span), help: Some("You should use `collect` to run your save command (see `help collect`). Or, you can put the file data in a variable and then pass the variable to `save`.".into()), inner: vec![], } } fn check_saving_to_source_file( metadata: Option<&PipelineMetadata>, dest: &Spanned, stderr_dest: Option<&Spanned>, ) -> Result<(), ShellError> { let Some(DataSource::FilePath(source)) = metadata.map(|meta| &meta.data_source) else { return Ok(()); }; if &dest.item == source { return Err(saving_to_source_file_error(dest)); } if let Some(dest) = stderr_dest { if &dest.item == source { return Err(saving_to_source_file_error(dest)); } } Ok(()) } /// 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(Value::coerce_into_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.coerce_into_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.item; if !(force || append) && path.exists() { Err(ShellError::GenericError { error: "Destination file already exists".into(), msg: format!( "Destination file '{}' already exists", path.to_string_lossy() ), span: Some(span), help: Some("you can use -f, --force to force overwriting the destination".into()), inner: vec![], }) } 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().append(true).open(path), _ => std::fs::File::create(path), }; file.map_err(|e| ShellError::GenericError { error: "Permission denied".into(), msg: e.to_string(), span: Some(span), help: None, inner: vec![], }) } /// Get output file and optional stderr file fn get_files( path: &Spanned, stderr_path: Option<&Spanned>, append: bool, out_append: bool, err_append: bool, force: bool, ) -> Result<(File, Option), ShellError> { // First check both paths let (path, path_span) = prepare_path(path, append || out_append, force)?; let stderr_path_and_span = stderr_path .as_ref() .map(|stderr_path| prepare_path(stderr_path, append || err_append, force)) .transpose()?; // Only if both files can be used open and possibly truncate them let file = open_file(path, path_span, append || out_append)?; let stderr_file = stderr_path_and_span .map(|(stderr_path, stderr_path_span)| { if path == stderr_path { Err(ShellError::GenericError { error: "input and stderr input to same file".into(), msg: "can't save both input and stderr input to the same file".into(), span: Some(stderr_path_span), help: Some("you should use `o+e> file` instead".into()), inner: vec![], }) } else { open_file(stderr_path, stderr_path_span, append || err_append) } }) .transpose()?; Ok((file, stderr_file)) } fn stream_to_file( mut stream: RawStream, mut file: File, span: Span, progress: bool, ) -> Result<(), ShellError> { // https://github.com/nushell/nushell/pull/9377 contains the reason // for not using BufWriter let writer = &mut 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) }; 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.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_all(&buf) { *process_failed_p = true; return Err(ShellError::IOError { msg: err.to_string(), }); } Ok(()) })?; // 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()); } } } file.flush()?; Ok(()) }