use crate::progress_bar; use nu_engine::get_eval_block; #[allow(deprecated)] use nu_engine::{command_prelude::*, current_dir}; use nu_path::expand_path_with; use nu_protocol::{ ast::{Expr, Expression}, byte_stream::copy_with_interrupt, process::ChildPipe, ByteStreamSource, DataSource, OutDest, PipelineMetadata, }; use std::{ fs::File, io::{self, BufRead, BufReader, Read, Write}, path::{Path, PathBuf}, sync::{atomic::AtomicBool, Arc}, 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::ByteStream(stream, 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, )?; let size = stream.known_size(); let ctrlc = engine_state.ctrlc.clone(); match stream.into_source() { ByteStreamSource::Read(read) => { stream_to_file(read, size, ctrlc, file, span, progress)?; } ByteStreamSource::File(source) => { stream_to_file(source, size, ctrlc, file, span, progress)?; } ByteStreamSource::Child(mut child) => { fn write_or_consume_stderr( stderr: ChildPipe, file: Option, span: Span, ctrlc: Option>, progress: bool, ) -> Result<(), ShellError> { if let Some(file) = file { match stderr { ChildPipe::Pipe(pipe) => { stream_to_file(pipe, None, ctrlc, file, span, progress) } ChildPipe::Tee(tee) => { stream_to_file(tee, None, ctrlc, file, span, progress) } }? } else { match stderr { ChildPipe::Pipe(mut pipe) => { io::copy(&mut pipe, &mut io::sink()) } ChildPipe::Tee(mut tee) => io::copy(&mut tee, &mut io::sink()), } .err_span(span)?; } Ok(()) } match (child.stdout.take(), child.stderr.take()) { (Some(stdout), stderr) => { // delegate a thread to redirect stderr to result. let handler = stderr .map(|stderr| { let ctrlc = ctrlc.clone(); thread::Builder::new().name("stderr saver".into()).spawn( move || { write_or_consume_stderr( stderr, stderr_file, span, ctrlc, progress, ) }, ) }) .transpose() .err_span(span)?; let res = match stdout { ChildPipe::Pipe(pipe) => { stream_to_file(pipe, None, ctrlc, file, span, progress) } ChildPipe::Tee(tee) => { stream_to_file(tee, None, ctrlc, 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)) => { write_or_consume_stderr( stderr, stderr_file, span, ctrlc, progress, )?; } (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 => { // It's not necessary to check if we are saving to the same file if this is a // collected value, and not a stream if !matches!(input, PipelineData::Value(..) | PipelineData::Empty) { 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( "insert a `collect` command in the pipeline before `save` (see `help collect`).".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 } else if let PipelineData::ByteStream(..) = input { None } else if let PipelineData::Value(Value::String { .. }, ..) = input { None } else { path.extension() .map(|name| name.to_string_lossy().to_string()) }; let input = if let Some(ext) = ext { convert_to_extension(engine_state, &ext, stack, input, span)? } else { input }; value_to_bytes(input.into_value(span)?) } /// 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 { if let Some(decl_id) = engine_state.find_decl(format!("to {extension}").as_bytes(), &[]) { let decl = engine_state.get_decl(decl_id); if let Some(block_id) = decl.block_id() { let block = engine_state.get_block(block_id); let eval_block = get_eval_block(engine_state); eval_block(engine_state, stack, block, input) } else { decl.run(engine_state, stack, &Call::new(span, span.past()), input) } } else { Ok(input) } } /// 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( source: impl Read, known_size: Option, ctrlc: Option>, mut file: File, span: Span, progress: bool, ) -> Result<(), ShellError> { // https://github.com/nushell/nushell/pull/9377 contains the reason for not using `BufWriter` if progress { let mut bytes_processed = 0; let mut bar = progress_bar::NuProgressBar::new(known_size); // TODO: reduce the number of progress bar updates? let mut reader = BufReader::new(source); let res = loop { if nu_utils::ctrl_c::was_pressed(&ctrlc) { bar.abandoned_msg("# Cancelled #".to_owned()); return Ok(()); } match reader.fill_buf() { Ok(&[]) => break Ok(()), Ok(buf) => { file.write_all(buf).err_span(span)?; let len = buf.len(); reader.consume(len); bytes_processed += len as u64; bar.update_bar(bytes_processed); } Err(e) if e.kind() == io::ErrorKind::Interrupted => continue, Err(e) => break Err(e), } }; // If the process failed, stop the progress bar with an error message. if let Err(err) = res { let _ = file.flush(); bar.abandoned_msg("# Error while saving #".to_owned()); Err(err.into_spanned(span).into()) } else { file.flush().err_span(span)?; Ok(()) } } else { copy_with_interrupt(source, file, span, ctrlc.as_deref())?; Ok(()) } }