onde_tqdm/
lib.rs

1//! Rust implementation of the popular Python command line progress bar tool tqdm.
2//!
3//! > The name "tqdm" derives from the Arabic word taqaddum (تقدّم) which can mean
4//! > "progress", and is an abbreviation for "I love you so much" in Spanish
5//! > (te quiero demasiado). Instantly make your loops show a smart progress
6//! > meter - just wrap any iterable with tqdm(iterable), and you're done!
7//!
8
9use std::*;
10
11use std::io::Write;
12use std::ops::{Deref, DerefMut};
13use std::time::{Duration, SystemTime};
14
15extern crate anyhow;
16use anyhow::Result;
17
18extern crate crossterm;
19use crossterm::QueueableCommand;
20use crossterm::{cursor, terminal};
21
22extern crate once_cell;
23use once_cell::sync::Lazy;
24
25#[cfg(test)]
26mod test;
27
28pub mod style;
29pub use style::{Colour, Style};
30
31pub mod lib_async;
32pub use lib_async::tqdm_async;
33
34/// Manually refresh all bars.
35pub fn refresh() -> Result<()> {
36    let mut out = io::stderr();
37
38    if let Ok(tqdm) = BAR.lock() {
39        let (ncols, nrows) = size();
40
41        if tqdm.is_empty() {
42            return Ok(());
43        }
44
45        out.queue(cursor::Hide)?;
46        out.queue(cursor::MoveToColumn(0))?;
47
48        let time = SystemTime::now();
49
50        for info in tqdm.values().take(nrows - 1) {
51            let bar = format!("{:<1$}", info.format(time)?, ncols);
52            out.queue(crossterm::style::Print(bar))?;
53        }
54
55        let nbars = tqdm.len();
56        if nbars >= nrows {
57            out.queue(terminal::Clear(terminal::ClearType::FromCursorDown))?;
58            out.queue(crossterm::style::Print(" ... (more hidden) ..."))?;
59            out.queue(cursor::MoveToColumn(0))?;
60        }
61
62        if let Some(rows) = num::NonZeroUsize::new(nbars - 1) {
63            out.queue(cursor::MoveUp(rows.get() as u16))?;
64        }
65
66        out.queue(cursor::Show)?;
67    }
68
69    Ok(out.flush()?)
70}
71
72/* -------------------------------------------------------------------------- */
73/*                                    TQDM                                    */
74/* -------------------------------------------------------------------------- */
75
76fn create<T>(n: Option<usize>, iter: T) -> Tqdm<T> {
77    let id = ID.fetch_add(1, sync::atomic::Ordering::SeqCst);
78    if let Ok(mut tqdm) = BAR.lock() {
79        tqdm.insert(
80            id,
81            Info {
82                config: Config::default(),
83
84                it: 0,
85                its: None,
86                total: n,
87
88                t0: SystemTime::now(),
89                prev: time::UNIX_EPOCH,
90            },
91        );
92    }
93
94    if let Err(err) = refresh() {
95        eprintln!("{err}")
96    }
97
98    Tqdm {
99        iter,
100        id,
101
102        next: time::UNIX_EPOCH,
103        step: 0,
104
105        mininterval: Duration::from_secs_f64(1. / 24.),
106        miniters: 1,
107    }
108}
109
110/// Create a progress bar from an iterable.
111///
112/// This function creates a default progress bar object and registers it
113/// to the global collection. The returned iterator [Deref] to the given
114/// one and will update its tqdm whenever `next` is called.
115///
116/// ## Examples
117/// ```
118/// use tqdm::tqdm;
119///
120/// for i in tqdm(0..100) {
121///     /* Your loop logic here */
122/// }
123/// ```
124///
125pub fn tqdm<Iter: IntoIterator>(iterable: Iter) -> Tqdm<Iter::IntoIter> {
126    let iter = iterable.into_iter();
127    create(iter.size_hint().1, iter)
128}
129
130/// Manually create a progress bar.
131///
132///
133/// ## Examples
134/// ```
135/// use tqdm::pbar;
136/// let mut pbar = pbar(Some(44850));
137///
138/// for i in 0..300 {
139///     pbar.update(i).unwrap();
140///     /* Your loop logic here */
141/// }
142/// ```
143///
144pub fn pbar(total: Option<usize>) -> Tqdm<()> {
145    create(total, ())
146}
147
148/// Iterator wrapper that updates progress bar on `next`
149///
150///
151/// ## Examples
152///
153/// - Basic Usage
154/// ```
155/// for _ in tqdm(0..100) {
156///     thread::sleep(Duration::from_millis(10));
157/// }
158/// ```
159///
160/// - Composition
161/// ```
162/// for _ in tqdm(tqdm(0..100).take(50)) {
163///     thread::sleep(Duration::from_millis(10));
164/// }
165/// ```
166///
167/// - Multi-threading
168/// ```
169/// let threads: Vec<_> = [200, 400, 100].iter().map(|its| {
170///         std::thread::spawn(move || {
171///             for _ in tqdm(0..*its) {
172///                 thread::sleep(Duration::from_millis(10));
173///             }
174///         })
175///     })
176///     .collect();
177///
178/// for handle in threads {
179///     handle.join().unwrap();
180/// }
181/// ```
182pub struct Tqdm<T> {
183    pub iter: T,
184
185    /// Hash
186    id: usize,
187
188    /// Next refresh time
189    next: SystemTime,
190
191    /// Cached
192    step: usize,
193
194    /// Refresh limit
195    mininterval: Duration,
196    miniters: usize,
197}
198
199/// Builder patterns
200impl<T> Tqdm<T> {
201    /// Configure progress bar's name.
202    ///
203    /// * `desc` bar description
204    ///     - `Some(S)`: Named progress bar
205    ///     - `None`: Anonymous
206    ///
207    ///
208    /// ## Examples
209    /// ```
210    /// tqdm(0..100).desc(Some("Bar1"))
211    /// ```
212    ///
213    pub fn desc<S: ToString>(self, desc: Option<S>) -> Self {
214        self.set_desc(desc);
215        self
216    }
217
218    /// Configure progress bar's total.
219    ///
220    /// * `total` total number of items
221    ///     - `Some(n)`: Known length
222    ///     - `None`: Unknown length
223    ///
224    ///
225    /// ## Examples
226    /// ```
227    /// tqdm(0..100).total(Some(50))
228    /// ```
229    ///
230    pub fn total(self, total: Option<usize>) -> Self {
231        if let Ok(mut tqdm) = BAR.lock() {
232            let info = tqdm.get_mut(&self.id);
233            if let Some(info) = info {
234                info.total = total;
235            }
236        }
237
238        self
239    }
240
241    /// Configure progress bar's width.
242    ///
243    /// * `width` width limitation
244    ///     - `Some(usize)`: Fixed width regardless of terminal size
245    ///     - `None`: Expand to formatter limit or full terminal width
246    ///
247    ///
248    /// ## Examples
249    /// ```
250    /// tqdm(0..100).width(Some(100))
251    /// ```
252    ///
253    pub fn width(self, width: Option<usize>) -> Self {
254        if let Ok(mut tqdm) = BAR.lock() {
255            let info = tqdm.get_mut(&self.id);
256            if let Some(info) = info {
257                info.config.width = width;
258            }
259        }
260
261        self
262    }
263
264    /// Configure progress bar's style.
265    ///
266    /// * `style` bar style enum
267    ///
268    ///
269    /// ## Examples
270    /// ```
271    /// tqdm(0..100).style(tqdm::Style::Balloon)
272    /// ```
273    ///
274    pub fn style(self, style: Style) -> Self {
275        if let Ok(mut tqdm) = BAR.lock() {
276            let info = tqdm.get_mut(&self.id);
277            if let Some(info) = info {
278                info.config.style = style;
279            }
280        }
281
282        self
283    }
284
285    /// Configure progress bar's units.
286    ///
287    /// * `units` unit of measurement
288    ///
289    /// ## Examples
290    /// ```
291    /// tqdm(0..100).units("files")
292    /// ```
293    /// 
294    pub fn units<S: ToString>(self, units: S) -> Self {
295        if let Ok(mut tqdm) = BAR.lock() {
296            let info = tqdm.get_mut(&self.id);
297            if let Some(info) = info {
298                info.config.units = units.to_string();
299            }
300        }
301
302        self
303    }
304
305    /// Configure progress bar's color.
306    ///
307    /// * `colour` bar color enum
308    ///
309    /// ## Examples
310    /// ```
311    /// tqdm(0..100).colour(tqdm::Colour::Green)
312    /// ```
313    ///
314    pub fn colour(self, colour: Colour) -> Self {
315        if let Ok(mut tqdm) = BAR.lock() {
316            let info = tqdm.get_mut(&self.id);
317            if let Some(info) = info {
318                info.config.colour = colour;
319            }
320        }
321
322        self
323    }
324
325    /// Exponential smoothing factor.
326    ///
327    /// * `smoothing` weight for the current update
328    ///
329    ///
330    /// ## Examples
331    /// ```
332    /// tqdm(0..100).smoothing(0.9999)
333    /// ```
334    ///
335    pub fn smoothing(self, smoothing: f64) -> Self {
336        if let Ok(mut tqdm) = BAR.lock() {
337            let info = tqdm.get_mut(&self.id);
338            if let Some(info) = info {
339                info.config.smoothing = smoothing;
340            }
341        }
342
343        self
344    }
345
346    /// Behavior of after termination.
347    ///
348    /// * `clear` termination behavior
349    ///     - true: remove this bar as if never created
350    ///     - false: leave completed bar at the very top
351    ///
352    ///
353    /// ## Examples
354    /// ```
355    /// tqdm(0..100).clear(true)
356    /// ```
357    ///
358    pub fn clear(self, clear: bool) -> Self {
359        if let Ok(mut tqdm) = BAR.lock() {
360            let info = tqdm.get_mut(&self.id);
361            if let Some(info) = info {
362                info.config.clear = clear;
363            }
364        }
365
366        self
367    }
368}
369
370impl<T> Tqdm<T> {
371    /// Manually update the progress bar.
372    pub fn update(&mut self, n: usize) -> Result<()> {
373        self.step += n;
374
375        if self.step >= self.miniters {
376            let now = SystemTime::now();
377            if now >= self.next {
378                if let Ok(mut tqdm) = BAR.lock() {
379                    if let Some(info) = tqdm.get_mut(&self.id) {
380                        info.update(now, self.step);
381                        self.step = 0;
382                    }
383                }
384                refresh()?;
385
386                self.next = now + self.mininterval;
387            }
388        }
389
390        Ok(())
391    }
392
393    /// Set description of a progress bar.
394    pub fn set_desc<S: ToString>(&self, desc: Option<S>) {
395        if let Ok(mut tqdm) = BAR.lock() {
396            let info = tqdm.get_mut(&self.id);
397            if let Some(info) = info {
398                info.config.desc = desc.map(|desc| desc.to_string());
399            }
400        }
401    }
402
403    /// Manually close the bar and unregister it.
404    pub fn close(&mut self) -> Result<()> {
405        let time = SystemTime::now();
406        let mut out = io::stderr();
407
408        if let Ok(mut tqdm) = BAR.lock() {
409            if let Some(mut info) = tqdm.remove(&self.id) {
410                info.update(time, self.step);
411
412                out.queue(cursor::MoveToColumn(0))?;
413
414                if info.config.clear {
415                    out.queue(cursor::MoveDown(tqdm.len() as u16))?;
416                    out.queue(terminal::Clear(terminal::ClearType::CurrentLine))?;
417                    out.queue(cursor::MoveUp(tqdm.len() as u16))?;
418                } else {
419                    out.queue(crossterm::style::Print(info.format(time)?))?;
420                    out.queue(crossterm::style::Print("\n"))?;
421                }
422            }
423        }
424
425        refresh()
426    }
427}
428
429impl<Iter: Iterator> Iterator for Tqdm<Iter> {
430    type Item = Iter::Item;
431
432    fn next(&mut self) -> Option<Self::Item> {
433        if let Some(next) = self.iter.next() {
434            if let Err(err) = self.update(1) {
435                eprintln!("{err}");
436            }
437            Some(next)
438        } else {
439            None
440        }
441    }
442
443    fn size_hint(&self) -> (usize, Option<usize>) {
444        self.iter.size_hint()
445    }
446}
447
448impl<Iter: Iterator> Deref for Tqdm<Iter> {
449    type Target = Iter;
450
451    fn deref(&self) -> &Self::Target {
452        &self.iter
453    }
454}
455
456impl<Iter: Iterator> DerefMut for Tqdm<Iter> {
457    fn deref_mut(&mut self) -> &mut Self::Target {
458        &mut self.iter
459    }
460}
461
462impl<T> Drop for Tqdm<T> {
463    fn drop(&mut self) {
464        if let Err(err) = self.close() {
465            eprintln!("{err}")
466        }
467    }
468}
469
470/* -------------------------------------------------------------------------- */
471/*                                    TRAIT                                   */
472/* -------------------------------------------------------------------------- */
473
474/// Trait that allows calling `.tqdm()`, equivalent to `tqdm::tqdm(iter)`.
475///
476///
477/// ## Examples
478/// ```
479/// use tqdm::Iter;
480/// (0..).take(1000).tqdm()
481/// ```
482///
483pub trait Iter<Item>: Iterator<Item = Item> {
484    fn tqdm(self) -> Tqdm<Self>
485    where
486        Self: Sized,
487    {
488        tqdm(self)
489    }
490}
491
492impl<Iter: Iterator> crate::Iter<Iter::Item> for Iter {}
493
494/* -------------------------------------------------------------------------- */
495/*                                   PRIVATE                                  */
496/* -------------------------------------------------------------------------- */
497
498/* --------------------------------- STATIC --------------------------------- */
499
500static ID: sync::atomic::AtomicUsize = sync::atomic::AtomicUsize::new(0);
501static BAR: Lazy<sync::Mutex<collections::BTreeMap<usize, Info>>> =
502    Lazy::new(|| sync::Mutex::new(collections::BTreeMap::new()));
503
504fn size<T: From<u16>>() -> (T, T) {
505    let (width, height) = terminal::size().unwrap_or((80, 24));
506    (T::from(width), T::from(height))
507}
508
509fn ftime(seconds: usize) -> String {
510    let m = seconds / 60 % 60;
511    let s = seconds % 60;
512    match seconds / 3600 {
513        0 => format!("{m:02}:{s:02}"),
514        h => format!("{h:02}:{m:02}:{s:02}"),
515    }
516}
517
518/* --------------------------------- CONFIG --------------------------------- */
519
520struct Config {
521    desc: Option<String>,
522    width: Option<usize>,
523    style: style::Style,
524    units: String,
525    colour: style::Colour,
526    smoothing: f64,
527    clear: bool,
528}
529
530impl Default for Config {
531    fn default() -> Self {
532        Config {
533            desc: None,
534            width: None,
535            style: Style::default(),
536            units: String::from("it"),
537            colour: Colour::default(),
538            smoothing: 0.3,
539            clear: false,
540        }
541    }
542}
543
544/* ---------------------------------- INFO ---------------------------------- */
545
546struct Info {
547    config: Config,
548
549    it: usize,
550    its: Option<f64>,
551    total: Option<usize>,
552
553    t0: SystemTime,
554    prev: SystemTime,
555}
556
557impl Info {
558    fn format(&self, t: SystemTime) -> Result<String> {
559        let desc = match &self.config.desc {
560            Some(s) => s.to_owned() + ": ",
561            None => String::new(),
562        };
563
564        let units = self.config.units.deref();
565
566        let elapsed = ftime(t.duration_since(self.t0)?.as_secs_f64() as usize);
567        let width = self.config.width.unwrap_or_else(|| size().0);
568
569        let it = self.it;
570        let its = match self.its {
571            None => String::from("?"),
572            Some(its) => format!("{its:.02}"),
573        };
574
575        Ok(match self.total.filter(|&total| total >= it) {
576            None => format_args!("{desc}{it}{units} [{elapsed}, {its}{units}/s]").to_string(),
577
578            Some(total) => {
579                let pct = (it as f64 / total as f64).clamp(0.0, 1.0);
580                let eta = match self.its {
581                    None => String::from("?"),
582                    Some(its) => ftime(((total - it) as f64 / its) as usize),
583                };
584
585                let colour_code = self.config.colour.ansi_code();
586                let reset_code = if colour_code.is_empty() {
587                    ""
588                } else {
589                    Colour::reset()
590                };
591
592                let bra_ = format!("{desc}{:>3}%|", (100.0 * pct) as usize);
593                let _ket = format!("| {it}/{total} [{elapsed}<{eta}, {its}i{units}/s]");
594                let tqdm = {
595                    if let Style::Pacman = self.config.style {
596                        let limit = (width.saturating_sub(bra_.len() + _ket.len()) / 3) * 3 - 6;
597                        let pattern: Vec<_> = self.config.style.to_string().chars().collect();
598
599                        let m = pattern.len();
600                        let n = ((limit as f64 * pct) * m as f64) as usize;
601
602                        let bar = pattern.last().unwrap().to_string().repeat(n / m);
603                        let empty = " o ".repeat(limit / 3 + 2)[bar.len() + 1..].to_string();
604
605                        match n / m {
606                            x if x == limit => bar,
607                            _ => format!("{bar}{}{empty}", pattern[0]),
608                        }
609                    } else {
610                        let limit = width.saturating_sub(bra_.len() + _ket.len());
611                        let pattern: Vec<_> = self.config.style.to_string().chars().collect();
612
613                        let m = pattern.len();
614                        let n = ((limit as f64 * pct) * m as f64) as usize;
615
616                        let bar = pattern.last().unwrap().to_string().repeat(n / m);
617                        match n / m {
618                            x if x == limit => bar,
619                            _ => format!("{:<limit$}", format!("{}{}", bar, pattern[n % m])),
620                        }
621                    }
622                };
623
624                format_args!("{bra_}{colour_code}{tqdm}{reset_code}{_ket}").to_string()
625            }
626        })
627    }
628
629    fn update(&mut self, t: SystemTime, n: usize) {
630        if self.prev != time::UNIX_EPOCH {
631            let dt = t.duration_since(self.prev).unwrap();
632            let its = n as f64 / dt.as_secs_f64();
633
634            self.its = match self.its {
635                None => Some(its),
636                Some(ema) => {
637                    let beta = self.config.smoothing;
638                    Some(its * beta + ema * (1. - beta))
639                }
640            };
641        }
642
643        self.prev = t;
644        self.it += n;
645    }
646}
onde_tqdm/lib.rs

onde_tqdm/
lib.rs
