simple_tqdm/
lib.rs

1//! simple-tqdm is a small wrapper around indicatif that tries to be similar to python's [`tqdm`](https://pypi.org/project/tqdm/) library.
2//!
3//! tqdm comes with both a `tqdm` function and a
4//!`Tqdm` trait depending on your preference.
5//!
6//! # Example
7//! ```
8//! use simple_tqdm::tqdm;
9//!
10//! for _ in tqdm(0..2 << 24) {}
11//! ```
12//! <img src="https://github.com/Icerath/simple-tqdm/blob/main/screenshots/large.gif?raw=true">
13//! Or if you'd like to customize the progress bar:
14//!
15//! ```
16//! use simple_tqdm::{Tqdm, Config};
17//!
18//! let config = Config::new().with_unit("num");
19//! for _ in (0..2 << 24).tqdm_config(config) {}
20//! ```
21//! Or if you'd like multiple bars.
22//!
23//! ```
24//! use simple_tqdm::{Tqdm, Config};
25//! fn main() {
26//!     let config = Config::new().with_progress_chars(">= ");
27//!     std::thread::scope(|scope| {
28//!         for _ in 0..3 {
29//!             scope.spawn(|| for _ in (0..2 << 24).tqdm_config(config.clone()) {});
30//!         }   
31//!     });
32//! }
33//! ```
34//! <img src="https://github.com/Icerath/simple-tqdm/blob/main/screenshots/multiple.gif?raw=true">
35//!
36//! # Parallel Iterators
37//! tqdm also has optional support for parallel
38//! iterators with [Rayon](https://github.com/rayon-rs/rayon). In your
39//! `Cargo.toml`, use the "rayon" feature:
40//!
41//! ```toml
42//! [dependencies]
43//! simple-tqdm = {version = "*", features = ["rayon"]}
44//! ```
45//! And then use it like this:
46//! ```
47//! use simple_tqdm::ParTqdm;
48//! use rayon::prelude::*;
49//!
50//! let vec: Vec<_> = (0..100000).into_par_iter().tqdm().map(|i| i + 1).collect();
51//! assert_eq!(vec[0], 1);
52//! ```
53
54mod config;
55
56pub use config::Config;
57use indicatif::{
58    MultiProgress, ProgressBar, ProgressBarIter, ProgressDrawTarget, ProgressState, ProgressStyle,
59};
60use std::{borrow::Cow, fmt::Write, sync::OnceLock};
61
62#[cfg(feature = "rayon")]
63mod parallel;
64#[cfg(feature = "rayon")]
65pub use parallel::{par_tqdm, ParTqdm};
66
67/// Wraps an iterator to display it's progress.
68pub trait Tqdm<I>: Sized {
69    /// Wraps an iterator to display it's progress.
70    /// Allows configuration of the progress bar.
71    fn tqdm_config(self, config: Config) -> ProgressBarIter<I>;
72    /// Wraps an iterator to display it's progress.
73    fn tqdm(self) -> ProgressBarIter<I> {
74        self.tqdm_config(Config::default())
75    }
76}
77
78static BARS: OnceLock<MultiProgress> = OnceLock::new();
79
80impl<I: Iterator> Tqdm<I> for I {
81    fn tqdm_config(self, config: Config) -> ProgressBarIter<I> {
82        let bars = BARS.get_or_init(MultiProgress::new);
83        bars.add(progress_bar(config, self.size_hint().1)).wrap_iter(self)
84    }
85}
86
87/// Equivalent to using the Tqdm trait method.
88pub fn tqdm<I: ExactSizeIterator>(iter: I) -> ProgressBarIter<I> {
89    iter.tqdm()
90}
91
92fn progress_bar(config: Config, iter_len: Option<usize>) -> ProgressBar {
93    let len = match config.total {
94        Some(total) => Some((total * config.unit_scale) as u64),
95        None => iter_len.map(|int| int as u64),
96    };
97
98    let draw_target =
99        if config.disable { ProgressDrawTarget::hidden() } else { ProgressDrawTarget::stderr() };
100
101    ProgressBar::with_draw_target(len, draw_target)
102        .with_finish(config.progress_finish())
103        .with_prefix(config.prefix)
104        .with_style(style(
105            config.unit,
106            config.unit_scale,
107            config.postfix,
108            &config.progress_chars,
109            &config.colour,
110        ))
111}
112
113#[allow(clippy::float_cmp)]
114fn style(
115    unit: Cow<'static, str>,
116    unit_scale: f64,
117    postfix: Cow<'static, str>,
118    progress_chars: &str,
119    color: &str,
120) -> ProgressStyle {
121    ProgressStyle::with_template(
122        &format!("{{prefix}}{{percent}}|{{wide_bar:.{color}}}| {{pos}}/{{len}} [{{elapsed}}<{{eta}}, {{per_sec}}{{postfix}}]",
123    ))
124    .unwrap()
125    .with_key(
126        "per_sec",
127        move |state: &ProgressState, w: &mut dyn Write| {
128            let _ = write!(w, "{:.2}{}/s", unit_scale * state.per_sec(), unit);
129        },
130    )
131    .with_key("percent", |state: &ProgressState, w: &mut dyn Write| {
132        let _ = write!(w, "{: >3}%", (state.fraction() * 100.0) as i32);
133    })
134    .with_key("elapsed", |state: &ProgressState, w: &mut dyn Write| {
135        let duration = state.elapsed();
136        let minutes = duration.as_secs() / 60;
137        let seconds = duration.as_secs() % 60;
138        let _ = write!(w, "{minutes:0>2}:{seconds:0>2}");
139    })
140    .with_key("eta", |state: &ProgressState, w: &mut dyn Write| {
141        let duration = state.eta();
142        let minutes = duration.as_secs() / 60;
143        let seconds = duration.as_secs() % 60;
144        let _ = write!(w, "{minutes:0>2}:{seconds:0>2}");
145    })
146    .with_key("pos", move |state: &ProgressState, w: &mut dyn Write| {
147        if unit_scale.round() == unit_scale {
148            let _ = write!(w, "{:?}", unit_scale as i64 * state.pos() as i64);
149        } else {
150            let _ = write!(w, "{:?}", unit_scale * state.pos() as f64);
151        }
152    })
153    .with_key("len", move |state: &ProgressState, w: &mut dyn Write| {
154        if unit_scale.round() == unit_scale {
155            let state_len = state.len().unwrap_or(state.pos()) as i64;
156            let _ = write!(w, "{:?}", unit_scale as i64 * state_len);
157        } else {
158            let state_len = state.len().unwrap_or(state.pos()) as f64;
159            let _ = write!(w, "{:?}", unit_scale * state_len);
160        }
161    })
162    .with_key("postfix", move |_: &ProgressState, w: &mut dyn Write| {
163        let _ = write!(w, "{postfix}");
164    })
165    .progress_chars(progress_chars)
166}
simple_tqdm/lib.rs

simple_tqdm/
lib.rs
