src/progress.rs

   1 //! Module that provides functions to display progress bars for compressing and decompressing files.
   2 use std::{
   3     io,
   4     sync::mpsc::{self, Receiver, Sender},
   5     thread,
   6     time::Duration,
   7 };
   8
   9 use indicatif::{ProgressBar, ProgressStyle};
  10
  11 /// Draw a ProgressBar using a function that checks periodically for the progress
  12 pub struct Progress {
  13     draw_stop: Sender<()>,
  14     clean_done: Receiver<()>,
  15     display_handle: DisplayHandle,
  16 }
  17
  18 /// Writes to this struct will be displayed on the progress bar or stdout depending on the
  19 /// ProgressBarPolicy
  20 struct DisplayHandle {
  21     buf: Vec<u8>,
  22     sender: Sender<String>,
  23 }
  24 impl io::Write for DisplayHandle {
  25     fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
  26         self.buf.extend(buf);
  27         // Newline is the signal to flush
  28         if matches!(buf.last(), Some(&b'\n')) {
  29             self.buf.pop();
  30             self.flush()?;
  31         }
  32         Ok(buf.len())
  33     }
  34
  35     fn flush(&mut self) -> io::Result<()> {
  36         fn io_error<X>(_: X) -> io::Error {
  37             io::Error::new(io::ErrorKind::Other, "failed to flush buffer")
  38         }
  39         self.sender
  40             .send(String::from_utf8(self.buf.drain(..).collect()).map_err(io_error)?)
  41             .map_err(io_error)
  42     }
  43 }
  44
  45 impl Progress {
  46     /// Create a ProgressBar using a function that checks periodically for the progress
  47     /// If precise is true, the total_input_size will be displayed as the total_bytes size
  48     /// If ACCESSIBLE is set, this function returns None
  49     pub fn new_accessible_aware(
  50         total_input_size: u64,
  51         precise: bool,
  52         current_position_fn: Option<Box<dyn Fn() -> u64 + Send>>,
  53     ) -> Option<Self> {
  54         if *crate::cli::ACCESSIBLE.get().unwrap() {
  55             return None;
  56         }
  57         Some(Self::new(total_input_size, precise, current_position_fn))
  58     }
  59
  60     fn new(total_input_size: u64, precise: bool, current_position_fn: Option<Box<dyn Fn() -> u64 + Send>>) -> Self {
  61         let (draw_tx, draw_rx) = mpsc::channel();
  62         let (clean_tx, clean_rx) = mpsc::channel();
  63         let (msg_tx, msg_rx) = mpsc::channel();
  64
  65         thread::spawn(move || {
  66             let template = {
  67                 let mut t = String::new();
  68                 t += "{wide_msg} [{elapsed_precise}] ";
  69                 if precise && current_position_fn.is_some() {
  70                     t += "[{bar:.cyan/blue}] ";
  71                 } else {
  72                     t += "{spinner:.green} ";
  73                 }
  74                 if current_position_fn.is_some() {
  75                     t += "{bytes}/ ";
  76                 }
  77                 if precise {
  78                     t += "{total_bytes} ";
  79                 }
  80                 t += "({bytes_per_sec}, {eta}) {path}";
  81                 t
  82             };
  83             let bar = ProgressBar::new(total_input_size);
  84             bar.set_style(ProgressStyle::default_bar().template(&template).progress_chars("#>-"));
  85
  86             while draw_rx.try_recv().is_err() {
  87                 if let Some(ref pos_fn) = current_position_fn {
  88                     bar.set_position(pos_fn());
  89                 } else {
  90                     bar.tick();
  91                 }
  92                 if let Ok(msg) = msg_rx.try_recv() {
  93                     bar.set_message(msg);
  94                 }
  95                 thread::sleep(Duration::from_millis(100));
  96             }
  97             bar.finish();
  98             let _ = clean_tx.send(());
  99         });
 100
 101         Progress {
 102             draw_stop: draw_tx,
 103             clean_done: clean_rx,
 104             display_handle: DisplayHandle {
 105                 buf: Vec::new(),
 106                 sender: msg_tx,
 107             },
 108         }
 109     }
 110
 111     pub(crate) fn display_handle(&mut self) -> &mut dyn io::Write {
 112         &mut self.display_handle
 113     }
 114 }
 115 impl Drop for Progress {
 116     fn drop(&mut self) {
 117         let _ = self.draw_stop.send(());
 118         let _ = self.clean_done.recv();
 119     }
 120 }