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.send(String::from_utf8(self.buf.drain(..).collect()).map_err(io_error)?).map_err(io_error)
  40     }
  41 }
  42
  43 impl Progress {
  44     /// Create a ProgressBar using a function that checks periodically for the progress
  45     /// If precise is true, the total_input_size will be displayed as the total_bytes size
  46     /// If ACCESSIBLE is set, this function returns None
  47     pub fn new_accessible_aware(
  48         total_input_size: u64,
  49         precise: bool,
  50         current_position_fn: Option<Box<dyn Fn() -> u64 + Send>>,
  51     ) -> Option<Self> {
  52         if *crate::cli::ACCESSIBLE.get().unwrap() {
  53             return None;
  54         }
  55         Some(Self::new(total_input_size, precise, current_position_fn))
  56     }
  57
  58     fn new(total_input_size: u64, precise: bool, current_position_fn: Option<Box<dyn Fn() -> u64 + Send>>) -> Self {
  59         let (draw_tx, draw_rx) = mpsc::channel();
  60         let (clean_tx, clean_rx) = mpsc::channel();
  61         let (msg_tx, msg_rx) = mpsc::channel();
  62
  63         thread::spawn(move || {
  64             let template = {
  65                 let mut t = String::new();
  66                 t += "{prefix} [{elapsed_precise}] ";
  67                 if precise && current_position_fn.is_some() {
  68                     t += "[{wide_bar:.cyan/blue}] ";
  69                 } else {
  70                     t += "{spinner:.green} ";
  71                 }
  72                 if current_position_fn.is_some() {
  73                     t += "{bytes}/ ";
  74                 }
  75                 if precise {
  76                     t += "{total_bytes} ";
  77                 }
  78                 t += "({bytes_per_sec}, {eta}) {path}";
  79                 t
  80             };
  81             let pb = ProgressBar::new(total_input_size);
  82             pb.set_style(ProgressStyle::default_bar().template(&template).progress_chars("#>-"));
  83
  84             while draw_rx.try_recv().is_err() {
  85                 if let Some(ref pos_fn) = current_position_fn {
  86                     pb.set_position(pos_fn());
  87                 } else {
  88                     pb.tick();
  89                 }
  90                 if let Ok(msg) = msg_rx.try_recv() {
  91                     pb.set_prefix(msg);
  92                 }
  93                 thread::sleep(Duration::from_millis(100));
  94             }
  95             pb.finish();
  96             let _ = clean_tx.send(());
  97         });
  98
  99         Progress {
 100             draw_stop: draw_tx,
 101             clean_done: clean_rx,
 102             display_handle: DisplayHandle { buf: Vec::new(), sender: msg_tx },
 103         }
 104     }
 105
 106     pub(crate) fn display_handle(&mut self) -> &mut dyn io::Write {
 107         &mut self.display_handle
 108     }
 109 }
 110 impl Drop for Progress {
 111     fn drop(&mut self) {
 112         let _ = self.draw_stop.send(());
 113         let _ = self.clean_done.recv();
 114     }
 115 }