nushell/crates/nu-command/src/filesystem/save.rs

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<PipelineData, ShellError> {
        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::<Spanned<PathBuf>>(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::<Spanned<PathBuf>>(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<File>,
                            span: Span,
                            ctrlc: Option<Arc<AtomicBool>>,
                            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<Example> {
        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<OutDest>, Option<OutDest>) {
        (Some(OutDest::Capture), Some(OutDest::Capture))
    }
}

fn saving_to_source_file_error(dest: &Spanned<PathBuf>) -> 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<PathBuf>,
    stderr_dest: Option<&Spanned<PathBuf>>,
) -> 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<Vec<u8>, 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<PipelineData, ShellError> {
    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<Vec<u8>, 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::<Result<Vec<String>, 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<PathBuf>,
    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<File, ShellError> {
    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<PathBuf>,
    stderr_path: Option<&Spanned<PathBuf>>,
    append: bool,
    out_append: bool,
    err_append: bool,
    force: bool,
) -> Result<(File, Option<File>), 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<u64>,
    ctrlc: Option<Arc<AtomicBool>>,
    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(())
    }
}