use indicatif::{ProgressBar, ProgressStyle};
use std::io::IsTerminal;
use std::sync::atomic::{AtomicU64, Ordering};
pub struct Progress {
total: u64,
done: AtomicU64,
mode: Mode,
}
enum Mode {
Bar(ProgressBar),
Plain,
Silent,
}
const LOG_INTERVAL: u64 = 25;
impl Progress {
pub fn new(total: u64, silent: bool) -> Self {
let mode = if silent {
Mode::Silent
} else if std::io::stderr().is_terminal() {
let bar = ProgressBar::new(total);
bar.set_style(
ProgressStyle::with_template("{bar:40} {percent}%")
.unwrap()
.progress_chars("=> "),
);
Mode::Bar(bar)
} else {
Mode::Plain
};
Progress { total, done: AtomicU64::new(0), mode }
}
pub fn tick(&self) {
self.tick_by(1);
}
pub fn tick_by(&self, n: u64) {
let before = self.done.fetch_add(n, Ordering::Relaxed);
let after = before + n;
match &self.mode {
Mode::Bar(bar) => bar.set_position(after),
Mode::Plain => {
if after / LOG_INTERVAL != before / LOG_INTERVAL || after == self.total {
eprintln!("{}/{} images processed", after, self.total);
}
}
Mode::Silent => {}
}
}
pub fn println(&self, msg: &str) {
match &self.mode {
Mode::Bar(bar) => bar.println(msg),
Mode::Plain | Mode::Silent => eprintln!("{msg}"),
}
}
pub fn finish(self) {
if let Mode::Bar(bar) = self.mode {
bar.finish_and_clear();
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn silent_mode_tick_does_not_panic() {
let p = Progress::new(10, true);
for _ in 0..10 {
p.tick();
}
p.finish();
}
#[test]
fn silent_mode_println_still_prints() {
let p = Progress::new(5, true);
p.println("an error message");
}
#[test]
fn zero_total_does_not_panic() {
let p = Progress::new(0, true);
p.tick();
p.finish();
}
#[test]
fn silent_mode_tick_by_does_not_panic() {
let p = Progress::new(100, true);
p.tick_by(40);
p.tick_by(60);
p.finish();
}
#[test]
fn concurrent_tick_from_multiple_threads_reaches_correct_total() {
use std::sync::Arc;
let progress = Arc::new(Progress::new(1000, true));
let handles: Vec<_> = (0..10)
.map(|_| {
let p = Arc::clone(&progress);
std::thread::spawn(move || {
for _ in 0..100 {
p.tick();
}
})
})
.collect();
for h in handles {
h.join().unwrap();
}
assert_eq!(progress.done.load(Ordering::Relaxed), 1000);
}
}