diff --git a/Cargo.lock b/Cargo.lock index c16bcf889c..260839e709 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1785,6 +1785,18 @@ dependencies = [ "serde", ] +[[package]] +name = "indicatif" +version = "0.17.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4295cbb7573c16d310e99e713cf9e75101eb190ab31fccd35f2d2691b4352b19" +dependencies = [ + "console", + "number_prefix", + "portable-atomic", + "unicode-width", +] + [[package]] name = "inotify" version = "0.7.1" @@ -2650,6 +2662,7 @@ dependencies = [ "htmlescape", "ical", "indexmap", + "indicatif", "is-root", "itertools", "libc", @@ -3111,6 +3124,12 @@ dependencies = [ "libc", ] +[[package]] +name = "number_prefix" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3" + [[package]] name = "objc" version = "0.2.7" @@ -3708,6 +3727,12 @@ dependencies = [ "nom", ] +[[package]] +name = "portable-atomic" +version = "0.3.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26f6a7b87c2e435a3241addceeeff740ff8b7e76b74c13bf9acb17fa454ea00b" + [[package]] name = "powierza-coefficient" version = "1.0.2" diff --git a/crates/nu-command/Cargo.toml b/crates/nu-command/Cargo.toml index fddeea23c8..749bfbfa0b 100644 --- a/crates/nu-command/Cargo.toml +++ b/crates/nu-command/Cargo.toml @@ -51,6 +51,7 @@ fs_extra = "1.2.0" htmlescape = "0.3.1" ical = "0.7.0" indexmap = { version="1.7", features=["serde-1"] } +indicatif = "0.17.2" Inflector = "0.11" is-root = "0.1.2" itertools = "0.10.0" diff --git a/crates/nu-command/src/filesystem/open.rs b/crates/nu-command/src/filesystem/open.rs index b94b4a25ae..8cb7e65ece 100644 --- a/crates/nu-command/src/filesystem/open.rs +++ b/crates/nu-command/src/filesystem/open.rs @@ -139,6 +139,7 @@ impl Command for Open { Box::new(BufferedReader { input: buf_reader }), ctrlc, call_span, + None, )), stderr: None, exit_code: None, diff --git a/crates/nu-command/src/filesystem/save.rs b/crates/nu-command/src/filesystem/save.rs index 8b3b7ddbd0..da85fcc601 100644 --- a/crates/nu-command/src/filesystem/save.rs +++ b/crates/nu-command/src/filesystem/save.rs @@ -9,6 +9,8 @@ use std::fs::File; use std::io::{BufWriter, Write}; use std::path::Path; +use crate::progress_bar; + #[derive(Clone)] pub struct Save; @@ -47,6 +49,7 @@ impl Command for Save { .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) } @@ -60,6 +63,7 @@ impl Command for Save { 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; @@ -81,16 +85,16 @@ impl Command for Save { // delegate a thread to redirect stderr to result. let handler = stderr.map(|stderr_stream| match stderr_file { - Some(stderr_file) => { - std::thread::spawn(move || stream_to_file(stderr_stream, stderr_file, span)) - } + Some(stderr_file) => std::thread::spawn(move || { + stream_to_file(stderr_stream, stderr_file, span, progress) + }), None => std::thread::spawn(move || { let _ = stderr_stream.into_bytes(); Ok(PipelineData::empty()) }), }); - let res = stream_to_file(stream, file, span); + let res = stream_to_file(stream, file, span, progress); if let Some(h) = handler { h.join().map_err(|err| { ShellError::ExternalCommand( @@ -332,10 +336,29 @@ fn stream_to_file( mut stream: RawStream, file: File, span: Span, + progress: bool, ) -> Result { let mut writer = BufWriter::new(file); - stream + 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 { @@ -353,13 +376,39 @@ fn stream_to_file( )); } }, - Err(err) => return Err(err), + 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()) + .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 } diff --git a/crates/nu-command/src/formats/to/text.rs b/crates/nu-command/src/formats/to/text.rs index c0e059eea5..605fa4e482 100644 --- a/crates/nu-command/src/formats/to/text.rs +++ b/crates/nu-command/src/formats/to/text.rs @@ -50,6 +50,7 @@ impl Command for ToText { }), engine_state.ctrlc.clone(), span, + None, )), stderr: None, exit_code: None, diff --git a/crates/nu-command/src/lib.rs b/crates/nu-command/src/lib.rs index f9f4e49825..ea3c8e8a99 100644 --- a/crates/nu-command/src/lib.rs +++ b/crates/nu-command/src/lib.rs @@ -20,6 +20,7 @@ mod misc; mod network; mod path; mod platform; +mod progress_bar; mod random; mod shells; mod sort_utils; diff --git a/crates/nu-command/src/network/fetch.rs b/crates/nu-command/src/network/fetch.rs index 969f4414a2..adf5f3194d 100644 --- a/crates/nu-command/src/network/fetch.rs +++ b/crates/nu-command/src/network/fetch.rs @@ -361,6 +361,26 @@ fn response_to_buffer( engine_state: &EngineState, span: Span, ) -> nu_protocol::PipelineData { + // Try to get the size of the file to be downloaded. + // This is helpful to show the progress of the stream. + let buffer_size = match &response.headers().get("content-length") { + Some(content_length) => { + let content_length = &(*content_length).clone(); // binding + + let content_length = content_length + .to_str() + .unwrap_or("") + .parse::() + .unwrap_or(0); + + if content_length == 0 { + None + } else { + Some(content_length) + } + } + _ => None, + }; let buffered_input = BufReader::new(response); PipelineData::ExternalStream { @@ -370,6 +390,7 @@ fn response_to_buffer( }), engine_state.ctrlc.clone(), span, + buffer_size, )), stderr: None, exit_code: None, diff --git a/crates/nu-command/src/network/post.rs b/crates/nu-command/src/network/post.rs index 5e70eb4e02..de94fb78a1 100644 --- a/crates/nu-command/src/network/post.rs +++ b/crates/nu-command/src/network/post.rs @@ -415,6 +415,7 @@ fn response_to_buffer( }), engine_state.ctrlc.clone(), span, + None, )), stderr: None, exit_code: None, diff --git a/crates/nu-command/src/progress_bar.rs b/crates/nu-command/src/progress_bar.rs new file mode 100644 index 0000000000..cb3ca3d0a8 --- /dev/null +++ b/crates/nu-command/src/progress_bar.rs @@ -0,0 +1,71 @@ +use indicatif::{ProgressBar, ProgressState, ProgressStyle}; +use std::fmt; + +// This module includes the progress bar used to show the progress when using the command `save` +// Eventually it would be nice to find a better place for it. + +pub struct NuProgressBar { + pub pb: ProgressBar, + bytes_processed: u64, + total_bytes: Option, +} + +impl NuProgressBar { + pub fn new(total_bytes: Option) -> NuProgressBar { + // Let's create the progress bar template. + let template = match total_bytes { + Some(_) => { + // We will use a progress bar if we know the total bytes of the stream + ProgressStyle::with_template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} {binary_bytes_per_sec} ({eta}) {msg}") + } + _ => { + // But if we don't know the total then we just show the stats progress + ProgressStyle::with_template( + "{spinner:.green} [{elapsed_precise}] {bytes} {binary_bytes_per_sec} {msg}", + ) + } + }; + + let total_bytes = match total_bytes { + Some(total_size) => total_size, + _ => 0, + }; + + let new_progress_bar = ProgressBar::new(total_bytes); + new_progress_bar.set_style( + template + .unwrap_or_else(|_| ProgressStyle::default_bar()) + .with_key("eta", |state: &ProgressState, w: &mut dyn fmt::Write| { + let _ = fmt::write(w, format_args!("{:.1}s", state.eta().as_secs_f64())); + }) + .progress_chars("#>-"), + ); + + NuProgressBar { + pb: new_progress_bar, + total_bytes: None, + bytes_processed: 0, + } + } + + pub fn update_bar(&mut self, bytes_processed: u64) { + self.pb.set_position(bytes_processed); + } + + // Commenting this for now but adding it in the future + //pub fn finished_msg(&self, msg: String) { + // self.pb.finish_with_message(msg); + //} + + pub fn abandoned_msg(&self, msg: String) { + self.pb.abandon_with_message(msg); + } + + pub fn clone(&self) -> NuProgressBar { + NuProgressBar { + pb: self.pb.clone(), + bytes_processed: self.bytes_processed, + total_bytes: self.total_bytes, + } + } +} diff --git a/crates/nu-command/src/system/run_external.rs b/crates/nu-command/src/system/run_external.rs index 6ecbcf28cb..6d0776a162 100644 --- a/crates/nu-command/src/system/run_external.rs +++ b/crates/nu-command/src/system/run_external.rs @@ -476,6 +476,7 @@ impl ExternalCommand { Box::new(stdout_receiver), output_ctrlc.clone(), head, + None, )) } else { None @@ -485,6 +486,7 @@ impl ExternalCommand { Box::new(stderr_receiver), output_ctrlc.clone(), head, + None, )) } else { None diff --git a/crates/nu-command/src/viewers/table.rs b/crates/nu-command/src/viewers/table.rs index 4bd07cdbae..6424e8c1c8 100644 --- a/crates/nu-command/src/viewers/table.rs +++ b/crates/nu-command/src/viewers/table.rs @@ -249,6 +249,7 @@ fn handle_table_command( ), ctrlc, call.head, + None, )), stderr: None, exit_code: None, @@ -732,6 +733,7 @@ fn handle_row_stream( }), ctrlc, head, + None, )), stderr: None, exit_code: None, diff --git a/crates/nu-protocol/src/pipeline_data.rs b/crates/nu-protocol/src/pipeline_data.rs index f526dd123a..939c3fb45c 100644 --- a/crates/nu-protocol/src/pipeline_data.rs +++ b/crates/nu-protocol/src/pipeline_data.rs @@ -566,6 +566,7 @@ impl PipelineData { Box::new(vec![Ok(stderr_bytes)].into_iter()), stderr_ctrlc, stderr_span, + None, ) }); diff --git a/crates/nu-protocol/src/value/stream.rs b/crates/nu-protocol/src/value/stream.rs index 1ee595fdfb..d9b3621888 100644 --- a/crates/nu-protocol/src/value/stream.rs +++ b/crates/nu-protocol/src/value/stream.rs @@ -10,6 +10,7 @@ pub struct RawStream { pub ctrlc: Option>, pub is_binary: bool, pub span: Span, + pub known_size: Option, // (bytes) } impl RawStream { @@ -17,6 +18,7 @@ impl RawStream { stream: Box, ShellError>> + Send + 'static>, ctrlc: Option>, span: Span, + known_size: Option, ) -> Self { Self { stream, @@ -24,6 +26,7 @@ impl RawStream { ctrlc, is_binary: false, span, + known_size, } } @@ -62,6 +65,7 @@ impl RawStream { ctrlc: self.ctrlc, is_binary: self.is_binary, span: self.span, + known_size: self.known_size, } } } diff --git a/src/main.rs b/src/main.rs index 7ad0fdf74f..fb9898cb2b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -316,6 +316,7 @@ fn main() -> Result<()> { Box::new(BufferedReader::new(buf_reader)), Some(ctrlc), redirect_stdin.span, + None, )), stderr: None, exit_code: None,