Only display progress bars when cli feature is enabled, allow passing progress callback (#283)

* Remove progress bars from flash target modules, make `indicatif` optional

* Make flash targets take an optional callback function to update progress

* Add back in our progress bars, but only when the `cli` feature is enabled

Author: jessebraham

espflash/Cargo.toml

@@ -47,7 +47,7 @@ directories-next = { version = "2.0.0", optional = true }
 esp-idf-part = "0.1.1"
 env_logger = { version = "0.9.1", optional = true }
 flate2 = "1.0.24"
-indicatif = "0.17.1"
+indicatif = { version = "0.17.1", optional = true }
 lazy_static = { version = "1.4.0", optional = true }
 log = "0.4.17"
 miette = { version = "5.3.0", features = ["fancy"] }
@@ -69,7 +69,7 @@ xmas-elf = "0.8.0"
 default = ["cli"]
 cli = [
     "dep:addr2line", "dep:clap", "dep:comfy-table", "dep:crossterm", "dep:dialoguer",
-    "dep:directories-next", "dep:env_logger", "dep:lazy_static", "dep:parse_int",
-    "dep:regex", "dep:serde-hex", "dep:update-informer"
+    "dep:directories-next", "dep:env_logger", "dep:indicatif", "dep:lazy_static",
+    "dep:parse_int", "dep:regex", "dep:serde-hex", "dep:update-informer"
 ]
 raspberry = ["dep:rppal"]

espflash/src/bin/espflash.rs

@@ -8,9 +8,10 @@ use std::{
 use clap::{Args, Parser, Subcommand};
 use espflash::{
     cli::{
-        self, board_info, clap_enum_variants, config::Config, connect, erase_partitions,
-        flash_elf_image, monitor::monitor, parse_partition_table, partition_table,
-        save_elf_as_image, serial_monitor, ConnectArgs, FlashConfigArgs, PartitionTableArgs,
+        self, board_info, build_progress_bar_callback, clap_enum_variants, config::Config, connect,
+        erase_partitions, flash_elf_image, monitor::monitor, parse_partition_table,
+        partition_table, progress_bar, save_elf_as_image, serial_monitor, ConnectArgs,
+        FlashConfigArgs, PartitionTableArgs,
     },
     image_format::ImageFormatKind,
     logging::initialize_logger,
@@ -206,7 +207,10 @@ fn write_bin(args: WriteBinArgs, config: &Config) -> Result<()> {
     let mut buffer = Vec::with_capacity(size.try_into().into_diagnostic()?);
     f.read_to_end(&mut buffer).into_diagnostic()?;
 
-    flasher.write_bin_to_flash(args.addr, &buffer)?;
+    let progress = progress_bar(format!("segment 0x{:X}", args.addr), None);
+    let progress_cb = Some(build_progress_bar_callback(progress));
+
+    flasher.write_bin_to_flash(args.addr, &buffer, progress_cb)?;
 
     Ok(())
 }

espflash/src/cli/mod.rs

@@ -8,6 +8,7 @@
 //! [espflash]: https://crates.io/crates/espflash
 
 use std::{
+    borrow::Cow,
     collections::HashMap,
     fs,
     io::Write,
@@ -17,7 +18,7 @@ use std::{
 use clap::{builder::ArgPredicate, Args};
 use comfy_table::{modifiers, presets::UTF8_FULL, Attribute, Cell, Color, Table};
 use esp_idf_part::{DataType, Partition, PartitionTable};
-use indicatif::HumanCount;
+use indicatif::{style::ProgressStyle, HumanCount, ProgressBar, ProgressDrawTarget};
 use log::{debug, info};
 use miette::{IntoDiagnostic, Result, WrapErr};
 use serialport::{SerialPortType, UsbPortInfo};
@@ -169,6 +170,57 @@ pub struct SaveImageArgs {
     pub skip_padding: bool,
 }
 
+/// Create a new [ProgressBar] with some message and styling applied
+pub fn progress_bar<S>(msg: S, len: Option<u64>) -> ProgressBar
+where
+    S: Into<Cow<'static, str>>,
+{
+    // If no length was provided, or the provided length is 0, we will initially
+    // hide the progress bar. This is done to avoid drawing a full-width progress
+    // bar before we've determined the length.
+    let draw_target = match len {
+        Some(len) if len > 0 => ProgressDrawTarget::stderr(),
+        _ => ProgressDrawTarget::hidden(),
+    };
+
+    let progress = ProgressBar::with_draw_target(len, draw_target)
+        .with_message(msg)
+        .with_style(
+            ProgressStyle::default_bar()
+                .template("[{elapsed_precise}] [{bar:40}] {pos:>7}/{len:7} {msg}")
+                .unwrap()
+                .progress_chars("=> "),
+        );
+
+    progress
+}
+
+/// Create a callback function for the provided [ProgressBar]
+pub fn build_progress_bar_callback(pb: ProgressBar) -> Box<dyn Fn(usize, usize)> {
+    // This is a bit of an odd callback function, as it handles the entire lifecycle
+    // of the progress bar.
+    Box::new(move |current: usize, total: usize| {
+        // If the length has not yet been set, set it and then change the draw target to
+        // make the progress bar visible.
+        match pb.length() {
+            Some(0) | None => {
+                pb.set_length(total as u64);
+                pb.set_draw_target(ProgressDrawTarget::stderr());
+            }
+            _ => {}
+        }
+
+        // Set the new position of the progress bar.
+        pb.set_position(current as u64);
+
+        // If we have reached the end, make sure to finish the progress bar to ensure
+        // proper output on the terminal.
+        if current == total {
+            pb.finish();
+        }
+    })
+}
+
 /// Select a serial port and establish a connection with a target device
 pub fn connect(args: &ConnectArgs, config: &Config) -> Result<Flasher> {
     let port_info = get_serial_port_info(args, config)?;

espflash/src/flasher/mod.rs

@@ -394,6 +394,7 @@ impl Flasher {
                     addr: text_addr,
                     data: Cow::Borrowed(&text),
                 },
+                None,
             )
             .flashing()?;
 
@@ -407,6 +408,7 @@ impl Flasher {
                     addr: data_addr,
                     data: Cow::Borrowed(&data),
                 },
+                None,
             )
             .flashing()?;
 
@@ -637,6 +639,9 @@ impl Flasher {
     /// Note that this will not touch the flash on the device
     pub fn load_elf_to_ram(&mut self, elf_data: &[u8]) -> Result<(), Error> {
         let image = ElfFirmwareImage::try_from(elf_data)?;
+        if image.rom_segments(self.chip).next().is_some() {
+            return Err(Error::ElfNotRamLoadable);
+        }
 
         let mut target = self.chip.ram_target(
             Some(image.entry()),
@@ -646,19 +651,21 @@ impl Flasher {
         );
         target.begin(&mut self.connection).flashing()?;
 
-        if image.rom_segments(self.chip).next().is_some() {
-            return Err(Error::ElfNotRamLoadable);
-        }
-
         for segment in image.ram_segments(self.chip) {
+            // Only display progress bars when the "cli" feature is enabled.
+            let progress_cb = if cfg!(feature = "cli") {
+                use crate::cli::{build_progress_bar_callback, progress_bar};
+
+                let progress = progress_bar(format!("segment 0x{:X}", segment.addr), None);
+                let progress_cb = build_progress_bar_callback(progress);
+
+                Some(progress_cb)
+            } else {
+                None
+            };
+
             target
-                .write_segment(
-                    &mut self.connection,
-                    RomSegment {
-                        addr: segment.addr,
-                        data: Cow::Borrowed(segment.data()),
-                    },
-                )
+                .write_segment(&mut self.connection, segment.into(), progress_cb)
                 .flashing()?;
         }
 
@@ -694,12 +701,25 @@ impl Flasher {
             flash_freq,
         )?;
 
+        // When the "cli" feature is enabled, display the image size information.
         #[cfg(feature = "cli")]
         crate::cli::display_image_size(image.app_size(), image.part_size());
 
         for segment in image.flash_segments() {
+            // Only display progress bars when the "cli" feature is enabled.
+            let progress_cb = if cfg!(feature = "cli") {
+                use crate::cli::{build_progress_bar_callback, progress_bar};
+
+                let progress = progress_bar(format!("segment 0x{:X}", segment.addr), None);
+                let progress_cb = build_progress_bar_callback(progress);
+
+                Some(progress_cb)
+            } else {
+                None
+            };
+
             target
-                .write_segment(&mut self.connection, segment)
+                .write_segment(&mut self.connection, segment, progress_cb)
                 .flashing()?;
         }
 
@@ -709,15 +729,22 @@ impl Flasher {
     }
 
     /// Load an bin image to flash at a specific address
-    pub fn write_bin_to_flash(&mut self, addr: u32, data: &[u8]) -> Result<(), Error> {
-        let mut target = self.chip.flash_target(self.spi_params, self.use_stub);
-        target.begin(&mut self.connection).flashing()?;
+    pub fn write_bin_to_flash(
+        &mut self,
+        addr: u32,
+        data: &[u8],
+        progress_cb: Option<Box<dyn Fn(usize, usize)>>,
+    ) -> Result<(), Error> {
         let segment = RomSegment {
             addr,
             data: Cow::from(data),
         };
-        target.write_segment(&mut self.connection, segment)?;
+
+        let mut target = self.chip.flash_target(self.spi_params, self.use_stub);
+        target.begin(&mut self.connection).flashing()?;
+        target.write_segment(&mut self.connection, segment, progress_cb)?;
         target.finish(&mut self.connection, true).flashing()?;
+
         Ok(())
     }
 

espflash/src/targets/flash_target/esp32.rs

@@ -4,7 +4,6 @@ use flate2::{
     write::{ZlibDecoder, ZlibEncoder},
     Compression,
 };
-use indicatif::{ProgressBar, ProgressStyle};
 
 use super::FlashTarget;
 use crate::{
@@ -36,49 +35,52 @@ impl Esp32Target {
 impl FlashTarget for Esp32Target {
     fn begin(&mut self, connection: &mut Connection) -> Result<(), Error> {
         connection.with_timeout(CommandType::SpiAttach.timeout(), |connection| {
-            connection.command(if self.use_stub {
+            let command = if self.use_stub {
                 Command::SpiAttachStub {
                     spi_params: self.spi_attach_params,
                 }
             } else {
                 Command::SpiAttach {
                     spi_params: self.spi_attach_params,
                 }
-            })
+            };
+
+            connection.command(command)
         })?;
 
-        // TODO remove this when we use the stub, the stub should be taking care of
-        // this. TODO do we also need to disable rtc super wdt?
+        // TODO: Remove this when we use the stub, the stub should be taking care of
+        //       this.
+        // TODO: Do we also need to disable RTC super WDT?
         if connection.get_usb_pid()? == USB_SERIAL_JTAG_PID {
             match self.chip {
                 Chip::Esp32c3 => {
                     connection.command(Command::WriteReg {
-                        address: 0x600080a8,
-                        value: 0x50D83AA1u32,
+                        address: 0x6000_80a8,
+                        value: 0x50D8_3AA1,
                         mask: None,
                     })?; // WP disable
                     connection.command(Command::WriteReg {
-                        address: 0x60008090,
+                        address: 0x6000_8090,
                         value: 0x0,
                         mask: None,
-                    })?; // turn off RTC WDG
+                    })?; // turn off RTC WDT
                     connection.command(Command::WriteReg {
-                        address: 0x600080a8,
+                        address: 0x6000_80a8,
                         value: 0x0,
                         mask: None,
                     })?; // WP enable
                 }
                 Chip::Esp32s3 => {
                     connection.command(Command::WriteReg {
                         address: 0x6000_80B0,
-                        value: 0x50D83AA1u32,
+                        value: 0x50D8_3AA1,
                         mask: None,
                     })?; // WP disable
                     connection.command(Command::WriteReg {
                         address: 0x6000_8098,
                         value: 0x0,
                         mask: None,
-                    })?; // turn off RTC WDG
+                    })?; // turn off RTC WDT
                     connection.command(Command::WriteReg {
                         address: 0x6000_80B0,
                         value: 0x0,
@@ -96,6 +98,7 @@ impl FlashTarget for Esp32Target {
         &mut self,
         connection: &mut Connection,
         segment: RomSegment,
+        progress_cb: Option<Box<dyn Fn(usize, usize)>>,
     ) -> Result<(), Error> {
         let addr = segment.addr;
 
@@ -126,16 +129,7 @@ impl FlashTarget for Esp32Target {
         )?;
 
         let chunks = compressed.chunks(flash_write_size);
-
-        let (_, chunk_size) = chunks.size_hint();
-        let chunk_size = chunk_size.unwrap_or(0) as u64;
-        let pb_chunk = ProgressBar::new(chunk_size);
-        pb_chunk.set_style(
-            ProgressStyle::default_bar()
-                .template("[{elapsed_precise}] [{bar:40}] {pos:>7}/{len:7} {msg}")
-                .unwrap()
-                .progress_chars("=> "),
-        );
+        let num_chunks = chunks.len();
 
         // decode the chunks to see how much data the device will have to save
         let mut decoder = ZlibDecoder::new(Vec::new());
@@ -147,7 +141,6 @@ impl FlashTarget for Esp32Target {
             let size = decoder.get_ref().len() - decoded_size;
             decoded_size = decoder.get_ref().len();
 
-            pb_chunk.set_message(format!("segment 0x{:X} writing chunks", addr));
             connection.with_timeout(
                 CommandType::FlashDeflateData.timeout_for_size(size as u32),
                 |connection| {
@@ -160,10 +153,11 @@ impl FlashTarget for Esp32Target {
                     Ok(())
                 },
             )?;
-            pb_chunk.inc(1);
-        }
 
-        pb_chunk.finish_with_message(format!("segment 0x{:X}", addr));
+            if let Some(ref cb) = progress_cb {
+                cb(i + 1, num_chunks);
+            }
+        }
 
         Ok(())
     }
@@ -172,10 +166,11 @@ impl FlashTarget for Esp32Target {
         connection.with_timeout(CommandType::FlashDeflateEnd.timeout(), |connection| {
             connection.command(Command::FlashDeflateEnd { reboot: false })
         })?;
+
         if reboot {
-            connection.reset()
-        } else {
-            Ok(())
+            connection.reset()?;
         }
+
+        Ok(())
     }
 }