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, byte_stream::copy_with_signals, process::ChildPipe, ByteStreamSource, DataSource, OutDest,
PipelineMetadata, Signals,
};
use std::{
fs::File,
io::{self, BufRead, BufReader, Read, Write},
path::{Path, PathBuf},
thread,
time::{Duration, Instant},
};
#[derive(Clone)]
pub struct Save;
impl Command for Save {
fn name(&self) -> &str {
"save"
}
fn description(&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 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, force)?;
let size = stream.known_size();
let signals = engine_state.signals();
match stream.into_source() {
ByteStreamSource::Read(read) => {
stream_to_file(read, size, signals, file, span, progress)?;
}
ByteStreamSource::File(source) => {
stream_to_file(source, size, signals, file, span, progress)?;
}
#[cfg(feature = "os")]
ByteStreamSource::Child(mut child) => {
fn write_or_consume_stderr(
stderr: ChildPipe,
file: Option<File>,
span: Span,
signals: &Signals,
progress: bool,
) -> Result<(), ShellError> {
if let Some(file) = file {
match stderr {
ChildPipe::Pipe(pipe) => {
stream_to_file(pipe, None, signals, file, span, progress)
}
ChildPipe::Tee(tee) => {
stream_to_file(tee, None, signals, file, span, progress)
}
}?
} else {
match stderr {
ChildPipe::Pipe(mut pipe) => {
io::copy(&mut pipe, &mut io::stderr())
}
ChildPipe::Tee(mut tee) => {
io::copy(&mut tee, &mut io::stderr())
}
}
.err_span(span)?;
}
Ok(())
}
match (child.stdout.take(), child.stderr.take()) {
(Some(stdout), stderr) => {
let handler = stderr
.map(|stderr| {
let signals = signals.clone();
thread::Builder::new().name("stderr saver".into()).spawn(
move || {
write_or_consume_stderr(
stderr,
stderr_file,
span,
&signals,
progress,
)
},
)
})
.transpose()
.err_span(span)?;
let res = match stdout {
ChildPipe::Pipe(pipe) => {
stream_to_file(pipe, None, signals, file, span, progress)
}
ChildPipe::Tee(tee) => {
stream_to_file(tee, None, signals, 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,
signals,
progress,
)?;
}
(None, None) => {}
};
child.wait()?;
}
}
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, 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 => {
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)?;
let (mut file, _) = get_files(&path, stderr_path.as_ref(), 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::PipeSeparate), Some(OutDest::PipeSeparate))
}
}
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(())
}
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)?)
}
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 {
let call = ast::Call::new(span);
decl.run(engine_state, stack, &(&call).into(), input)
}
} else {
Ok(input)
}
}
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())
}
Value::Error { error, .. } => Err(*error),
other => Ok(other.coerce_into_string()?.into_bytes()),
}
}
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: format!("Problem with [{}], Permission denied", path.display()),
msg: e.to_string(),
span: Some(span),
help: None,
inner: vec![],
})
}
fn get_files(
path: &Spanned<PathBuf>,
stderr_path: Option<&Spanned<PathBuf>>,
append: bool,
force: bool,
) -> Result<(File, Option<File>), ShellError> {
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()?;
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 {
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)
}
})
.transpose()?;
Ok((file, stderr_file))
}
fn stream_to_file(
source: impl Read,
known_size: Option<u64>,
signals: &Signals,
mut file: File,
span: Span,
progress: bool,
) -> Result<(), ShellError> {
if progress {
let mut bytes_processed = 0;
let mut bar = progress_bar::NuProgressBar::new(known_size);
let mut last_update = Instant::now();
let mut reader = BufReader::new(source);
let res = loop {
if let Err(err) = signals.check(span) {
bar.abandoned_msg("# Cancelled #".to_owned());
return Err(err);
}
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;
if last_update.elapsed() >= Duration::from_millis(75) {
bar.update_bar(bytes_processed);
last_update = Instant::now();
}
}
Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
Err(e) => break Err(e),
}
};
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_signals(source, file, span, signals)?;
Ok(())
}
}