Skip to main content

imdl_indicatif/
progress.rs

1use std::fmt;
2use std::io;
3use std::sync::atomic::{AtomicBool, Ordering};
4use std::sync::mpsc::{channel, Receiver, Sender};
5use std::sync::Arc;
6use std::sync::{Mutex, RwLock};
7use std::thread;
8use std::time::{Duration, Instant};
9
10use crate::style::ProgressStyle;
11use crate::utils::{duration_to_secs, secs_to_duration, Estimate};
12use console::Term;
13
14/// The drawn state of an element.
15#[derive(Clone, Debug)]
16struct ProgressDrawState {
17    /// The lines to print (can contain ANSI codes)
18    pub lines: Vec<String>,
19    /// The number of lines that shouldn't be reaped by the next tick.
20    pub orphan_lines: usize,
21    /// True if the bar no longer needs drawing.
22    pub finished: bool,
23    /// True if drawing should be forced.
24    pub force_draw: bool,
25    /// True if we should move the cursor up when possible instead of clearing lines.
26    pub move_cursor: bool,
27    /// Time when the draw state was created.
28    pub ts: Instant,
29}
30
31#[derive(Debug)]
32enum Status {
33    InProgress,
34    DoneVisible,
35    DoneHidden,
36}
37
38enum ProgressDrawTargetKind {
39    Term(Term, Option<ProgressDrawState>, Option<Duration>),
40    Remote(usize, Mutex<Sender<(usize, ProgressDrawState)>>),
41    Hidden,
42}
43
44/// Target for draw operations
45///
46/// This tells a progress bar or a multi progress object where to paint to.
47/// The draw target is a stateful wrapper over a drawing destination and
48/// internally optimizes how often the state is painted to the output
49/// device.
50pub struct ProgressDrawTarget {
51    kind: ProgressDrawTargetKind,
52}
53
54impl ProgressDrawTarget {
55    /// Draw to a buffered stdout terminal at a max of 15 times a second.
56    ///
57    /// For more information see `ProgressDrawTarget::to_term`.
58    pub fn stdout() -> ProgressDrawTarget {
59        ProgressDrawTarget::to_term(Term::buffered_stdout(), 15)
60    }
61
62    /// Draw to a buffered stderr terminal at a max of 15 times a second.
63    ///
64    /// This is the default draw target for progress bars.  For more
65    /// information see `ProgressDrawTarget::to_term`.
66    pub fn stderr() -> ProgressDrawTarget {
67        ProgressDrawTarget::to_term(Term::buffered_stderr(), 15)
68    }
69
70    /// Draw to a buffered stdout terminal at a max of `refresh_rate` times a second.
71    ///
72    /// For more information see `ProgressDrawTarget::to_term`.
73    pub fn stdout_with_hz(refresh_rate: u64) -> ProgressDrawTarget {
74        ProgressDrawTarget::to_term(Term::buffered_stdout(), refresh_rate)
75    }
76
77    /// Draw to a buffered stderr terminal at a max of `refresh_rate` times a second.
78    ///
79    /// For more information see `ProgressDrawTarget::to_term`.
80    pub fn stderr_with_hz(refresh_rate: u64) -> ProgressDrawTarget {
81        ProgressDrawTarget::to_term(Term::buffered_stderr(), refresh_rate)
82    }
83
84    /// Draw to a buffered stdout terminal without max framerate.
85    ///
86    /// This is useful when data is known to come in very slowly and
87    /// not rendering some updates would be a problem (for instance
88    /// when messages are used extensively).
89    ///
90    /// For more information see `ProgressDrawTarget::to_term`.
91    pub fn stdout_nohz() -> ProgressDrawTarget {
92        ProgressDrawTarget::to_term(Term::buffered_stdout(), None)
93    }
94
95    /// Draw to a buffered stderr terminal without max framerate.
96    ///
97    /// This is useful when data is known to come in very slowly and
98    /// not rendering some updates would be a problem (for instance
99    /// when messages are used extensively).
100    ///
101    /// For more information see `ProgressDrawTarget::to_term`.
102    pub fn stderr_nohz() -> ProgressDrawTarget {
103        ProgressDrawTarget::to_term(Term::buffered_stderr(), None)
104    }
105
106    /// Draw to a terminal, optionally with a specific refresh rate.
107    ///
108    /// Progress bars are by default drawn to terminals however if the
109    /// terminal is not user attended the entire progress bar will be
110    /// hidden.  This is done so that piping to a file will not produce
111    /// useless escape codes in that file.
112    pub fn to_term(term: Term, refresh_rate: impl Into<Option<u64>>) -> ProgressDrawTarget {
113        let rate = refresh_rate.into().map(|x| Duration::from_millis(1000 / x));
114        ProgressDrawTarget {
115            kind: ProgressDrawTargetKind::Term(term, None, rate),
116        }
117    }
118
119    /// A hidden draw target.
120    ///
121    /// This forces a progress bar to be not rendered at all.
122    pub fn hidden() -> ProgressDrawTarget {
123        ProgressDrawTarget {
124            kind: ProgressDrawTargetKind::Hidden,
125        }
126    }
127
128    /// Returns true if the draw target is hidden.
129    ///
130    /// This is internally used in progress bars to figure out if overhead
131    /// from drawing can be prevented.
132    pub fn is_hidden(&self) -> bool {
133        match self.kind {
134            ProgressDrawTargetKind::Hidden => true,
135            ProgressDrawTargetKind::Term(ref term, ..) => !term.is_term(),
136            _ => false,
137        }
138    }
139
140    /// Apply the given draw state (draws it).
141    fn apply_draw_state(&mut self, draw_state: ProgressDrawState) -> io::Result<()> {
142        // no need to apply anything to hidden draw targets.
143        if self.is_hidden() {
144            return Ok(());
145        }
146        match self.kind {
147            ProgressDrawTargetKind::Term(ref term, ref mut last_state, rate) => {
148                let last_draw = last_state.as_ref().map(|x| x.ts);
149                if draw_state.finished
150                    || draw_state.force_draw
151                    || rate.is_none()
152                    || last_draw.is_none()
153                    || last_draw.unwrap().elapsed() > rate.unwrap()
154                {
155                    if let Some(ref last_state) = *last_state {
156                        if !draw_state.lines.is_empty() && draw_state.move_cursor {
157                            last_state.move_cursor(term)?;
158                        } else {
159                            last_state.clear_term(term)?;
160                        }
161                    }
162                    draw_state.draw_to_term(term)?;
163                    term.flush()?;
164                    *last_state = Some(draw_state);
165                }
166            }
167            ProgressDrawTargetKind::Remote(idx, ref chan) => {
168                return chan
169                    .lock()
170                    .unwrap()
171                    .send((idx, draw_state))
172                    .map_err(|e| io::Error::new(io::ErrorKind::Other, e));
173            }
174            ProgressDrawTargetKind::Hidden => {}
175        }
176        Ok(())
177    }
178
179    /// Properly disconnects from the draw target
180    fn disconnect(&self) {
181        match self.kind {
182            ProgressDrawTargetKind::Term(_, _, _) => {}
183            ProgressDrawTargetKind::Remote(idx, ref chan) => {
184                chan.lock()
185                    .unwrap()
186                    .send((
187                        idx,
188                        ProgressDrawState {
189                            lines: vec![],
190                            orphan_lines: 0,
191                            finished: true,
192                            force_draw: false,
193                            move_cursor: false,
194                            ts: Instant::now(),
195                        },
196                    ))
197                    .ok();
198            }
199            ProgressDrawTargetKind::Hidden => {}
200        };
201    }
202}
203
204impl ProgressDrawState {
205    pub fn clear_term(&self, term: &Term) -> io::Result<()> {
206        term.clear_last_lines(self.lines.len() - self.orphan_lines)
207    }
208
209    pub fn move_cursor(&self, term: &Term) -> io::Result<()> {
210        term.move_cursor_up(self.lines.len() - self.orphan_lines)
211    }
212
213    pub fn draw_to_term(&self, term: &Term) -> io::Result<()> {
214        for line in &self.lines {
215            term.write_line(line)?;
216        }
217        Ok(())
218    }
219}
220
221/// The state of a progress bar at a moment in time.
222pub(crate) struct ProgressState {
223    pub(crate) style: ProgressStyle,
224    pub(crate) pos: u64,
225    pub(crate) len: u64,
226    pub(crate) tick: u64,
227    pub(crate) started: Instant,
228    draw_target: ProgressDrawTarget,
229    width: Option<u16>,
230    message: String,
231    prefix: String,
232    draw_delta: u64,
233    draw_next: u64,
234    status: Status,
235    est: Estimate,
236    tick_thread: Option<thread::JoinHandle<()>>,
237    steady_tick: u64,
238}
239
240impl ProgressState {
241    /// Returns the string that should be drawn for the
242    /// current spinner string.
243    pub fn current_tick_str(&self) -> &str {
244        if self.is_finished() {
245            self.style.get_final_tick_str()
246        } else {
247            self.style.get_tick_str(self.tick)
248        }
249    }
250
251    /// Indicates that the progress bar finished.
252    pub fn is_finished(&self) -> bool {
253        match self.status {
254            Status::InProgress => false,
255            Status::DoneVisible => true,
256            Status::DoneHidden => true,
257        }
258    }
259
260    /// Returns `false` if the progress bar should no longer be
261    /// drawn.
262    pub fn should_render(&self) -> bool {
263        match self.status {
264            Status::DoneHidden => false,
265            _ => true,
266        }
267    }
268
269    /// Returns the completion as a floating-point number between 0 and 1
270    pub fn fraction(&self) -> f32 {
271        let pct = match (self.pos, self.len) {
272            (_, 0) => 1.0,
273            (0, _) => 0.0,
274            (pos, len) => pos as f32 / len as f32,
275        };
276        pct.max(0.0).min(1.0)
277    }
278
279    /// Returns the position of the status bar as `(pos, len)` tuple.
280    pub fn position(&self) -> (u64, u64) {
281        (self.pos, self.len)
282    }
283
284    /// Returns the current message of the progress bar.
285    pub fn message(&self) -> &str {
286        &self.message
287    }
288
289    /// Returns the current prefix of the progress bar.
290    pub fn prefix(&self) -> &str {
291        &self.prefix
292    }
293
294    /// The entire draw width
295    pub fn width(&self) -> usize {
296        if let Some(width) = self.width {
297            width as usize
298        } else {
299            Term::stderr().size().1 as usize
300        }
301    }
302
303    /// Return the current average time per step
304    pub fn avg_time_per_step(&self) -> Duration {
305        self.est.time_per_step()
306    }
307
308    /// The expected ETA
309    pub fn eta(&self) -> Duration {
310        if self.len == !0 || self.is_finished() {
311            return Duration::new(0, 0);
312        }
313        let t = duration_to_secs(self.avg_time_per_step());
314        // add 0.75 to leave 0.25 sec of 0s for the user
315        secs_to_duration(t * self.len.saturating_sub(self.pos) as f64 + 0.75)
316    }
317
318    /// The number of steps per second
319    pub fn per_sec(&self) -> u64 {
320        let avg_time = self.avg_time_per_step().as_nanos();
321        if avg_time == 0 {
322            0
323        } else {
324            (1_000_000_000 / avg_time) as u64
325        }
326    }
327}
328
329/// A progress bar or spinner.
330///
331/// The progress bar is an `Arc` around an internal state.  When the progress
332/// bar is cloned it just increments the refcount which means the bar is
333/// shared with the original one.
334#[derive(Clone)]
335pub struct ProgressBar {
336    state: Arc<RwLock<ProgressState>>,
337}
338
339impl fmt::Debug for ProgressBar {
340    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
341        f.debug_struct("ProgressBar").finish()
342    }
343}
344
345impl ProgressBar {
346    /// Creates a new progress bar with a given length.
347    ///
348    /// This progress bar by default draws directly to stderr, and refreshes
349    /// a maximum of 15 times a second. To change the refresh rate set the
350    /// draw target to one with a different refresh rate.
351    pub fn new(len: u64) -> ProgressBar {
352        ProgressBar::with_draw_target(len, ProgressDrawTarget::stderr())
353    }
354
355    /// Creates a completely hidden progress bar.
356    ///
357    /// This progress bar still responds to API changes but it does not
358    /// have a length or render in any way.
359    pub fn hidden() -> ProgressBar {
360        ProgressBar::with_draw_target(!0, ProgressDrawTarget::hidden())
361    }
362
363    /// Creates a new progress bar with a given length and draw target.
364    pub fn with_draw_target(len: u64, target: ProgressDrawTarget) -> ProgressBar {
365        ProgressBar {
366            state: Arc::new(RwLock::new(ProgressState {
367                style: ProgressStyle::default_bar(),
368                draw_target: target,
369                width: None,
370                message: "".into(),
371                prefix: "".into(),
372                pos: 0,
373                len,
374                tick: 0,
375                draw_delta: 0,
376                draw_next: 0,
377                status: Status::InProgress,
378                started: Instant::now(),
379                est: Estimate::new(),
380                tick_thread: None,
381                steady_tick: 0,
382            })),
383        }
384    }
385
386    /// A convenience builder-like function for a progress bar with a given style.
387    pub fn with_style(self, style: ProgressStyle) -> ProgressBar {
388        self.state.write().unwrap().style = style;
389        self
390    }
391
392    /// Creates a new spinner.
393    ///
394    /// This spinner by default draws directly to stderr.  This adds the
395    /// default spinner style to it.
396    pub fn new_spinner() -> ProgressBar {
397        let rv = ProgressBar::new(!0);
398        rv.set_style(ProgressStyle::default_spinner());
399        rv
400    }
401
402    /// Overrides the stored style.
403    ///
404    /// This does not redraw the bar.  Call `tick` to force it.
405    pub fn set_style(&self, style: ProgressStyle) {
406        self.state.write().unwrap().style = style;
407    }
408
409    /// Spawns a background thread to tick the progress bar.
410    ///
411    /// When this is enabled a background thread will regularly tick the
412    /// progress back in the given interval (milliseconds).  This is
413    /// useful to advance progress bars that are very slow by themselves.
414    ///
415    /// When steady ticks are enabled calling `.tick()` on a progress
416    /// bar does not do anything.
417    pub fn enable_steady_tick(&self, ms: u64) {
418        let mut state = self.state.write().unwrap();
419        state.steady_tick = ms;
420        if state.tick_thread.is_some() {
421            return;
422        }
423
424        // Using a weak pointer is required to prevent a potential deadlock. See issue #133
425        let state_arc = Arc::downgrade(&self.state);
426        state.tick_thread = Some(thread::spawn(move || loop {
427            thread::sleep(Duration::from_millis(ms));
428            if let Some(state_arc) = state_arc.upgrade() {
429                let mut state = state_arc.write().unwrap();
430                if state.is_finished() || state.steady_tick == 0 {
431                    state.steady_tick = 0;
432                    state.tick_thread = None;
433                    break;
434                }
435                if state.tick != 0 {
436                    state.tick = state.tick.saturating_add(1);
437                }
438
439                draw_state(&mut state).ok();
440            } else {
441                break;
442            }
443        }));
444
445        // use the side effect of tick to force the bar to tick.
446        ::std::mem::drop(state);
447        self.tick();
448    }
449
450    /// Undoes `enable_steady_tick`.
451    pub fn disable_steady_tick(&self) {
452        self.enable_steady_tick(0);
453    }
454
455    /// Limit redrawing of progress bar to every `n` steps. Defaults to 0.
456    ///
457    /// By default, the progress bar will redraw whenever its state advances.
458    /// This setting is helpful in situations where the overhead of redrawing
459    /// the progress bar dominates the computation whose progress is being
460    /// reported.
461    ///
462    /// If `n` is greater than 0, operations that change the progress bar such
463    /// as `.tick()`, `.set_message()` and `.set_length()` will no longer cause
464    /// the progress bar to be redrawn, and will only be shown once the
465    /// position advances by `n` steps.
466    ///
467    /// ```rust,no_run
468    /// # use imdl_indicatif::ProgressBar;
469    /// let n = 1_000_000;
470    /// let pb = ProgressBar::new(n);
471    /// pb.set_draw_delta(n / 100); // redraw every 1% of additional progress
472    /// ```
473    ///
474    /// Note that `ProgressDrawTarget` may impose additional buffering of redraws.
475    pub fn set_draw_delta(&self, n: u64) {
476        let mut state = self.state.write().unwrap();
477        state.draw_delta = n;
478        state.draw_next = state.pos.saturating_add(state.draw_delta);
479    }
480
481    /// Manually ticks the spinner or progress bar.
482    ///
483    /// This automatically happens on any other change to a progress bar.
484    pub fn tick(&self) {
485        self.update_and_draw(|state| {
486            if state.steady_tick == 0 || state.tick == 0 {
487                state.tick = state.tick.saturating_add(1);
488            }
489        });
490    }
491
492    /// Advances the position of a progress bar by delta.
493    pub fn inc(&self, delta: u64) {
494        self.update_and_draw(|state| {
495            state.pos = state.pos.saturating_add(delta);
496            if state.steady_tick == 0 || state.tick == 0 {
497                state.tick = state.tick.saturating_add(1);
498            }
499        })
500    }
501
502    /// A quick convenience check if the progress bar is hidden.
503    pub fn is_hidden(&self) -> bool {
504        self.state.read().unwrap().draw_target.is_hidden()
505    }
506
507    // Indicates that the progress bar finished.
508    pub fn is_finished(&self) -> bool {
509        self.state.read().unwrap().is_finished()
510    }
511
512    /// Print a log line above the progress bar.
513    ///
514    /// If the progress bar was added to a `MultiProgress`, the log line will be
515    /// printed above all other progress bars.
516    ///
517    /// Note that if the progress bar is hidden (which by default happens if
518    /// the progress bar is redirected into a file) println will not do
519    /// anything either.
520    pub fn println<I: Into<String>>(&self, msg: I) {
521        let mut state = self.state.write().unwrap();
522
523        let mut lines: Vec<String> = msg.into().lines().map(Into::into).collect();
524        let orphan_lines = lines.len();
525        if state.should_render() {
526            lines.extend(state.style.format_state(&*state));
527        }
528
529        let draw_state = ProgressDrawState {
530            lines,
531            orphan_lines,
532            finished: state.is_finished(),
533            force_draw: true,
534            move_cursor: false,
535            ts: Instant::now(),
536        };
537
538        state.draw_target.apply_draw_state(draw_state).ok();
539    }
540
541    /// Sets the position of the progress bar.
542    pub fn set_position(&self, pos: u64) {
543        self.update_and_draw(|state| {
544            state.draw_next = pos;
545            state.pos = pos;
546            if state.steady_tick == 0 || state.tick == 0 {
547                state.tick = state.tick.saturating_add(1);
548            }
549        })
550    }
551
552    /// Sets the length of the progress bar.
553    pub fn set_length(&self, len: u64) {
554        self.update_and_draw(|state| {
555            state.len = len;
556        })
557    }
558
559    /// Increase the length of the progress bar.
560    pub fn inc_length(&self, delta: u64) {
561        self.update_and_draw(|state| {
562            state.len = state.len.saturating_add(delta);
563        })
564    }
565
566    /// Sets the current prefix of the progress bar.
567    ///
568    /// For the prefix to be visible, `{prefix}` placeholder
569    /// must be present in the template (see `ProgressStyle`).
570    pub fn set_prefix(&self, prefix: &str) {
571        let prefix = prefix.to_string();
572        self.update_and_draw(|state| {
573            state.prefix = prefix;
574            if state.steady_tick == 0 || state.tick == 0 {
575                state.tick = state.tick.saturating_add(1);
576            }
577        })
578    }
579
580    /// Sets the current message of the progress bar.
581    ///
582    /// For the message to be visible, `{msg}` placeholder
583    /// must be present in the template (see `ProgressStyle`).
584    pub fn set_message(&self, msg: &str) {
585        let msg = msg.to_string();
586        self.update_and_draw(|state| {
587            state.message = msg;
588            if state.steady_tick == 0 || state.tick == 0 {
589                state.tick = state.tick.saturating_add(1);
590            }
591        })
592    }
593
594    /// Resets the ETA calculation.
595    ///
596    /// This can be useful if progress bars make a huge jump or were
597    /// paused for a prolonged time.
598    pub fn reset_eta(&self) {
599        self.update_and_draw(|state| {
600            state.est.reset();
601        });
602    }
603
604    /// Resets elapsed time
605    pub fn reset_elapsed(&self) {
606        self.update_and_draw(|state| {
607            state.started = Instant::now();
608        });
609    }
610
611    pub fn reset(&self) {
612        self.reset_eta();
613        self.reset_elapsed();
614        self.update_and_draw(|state| {
615            state.draw_next = 0;
616            state.pos = 0;
617            state.status = Status::InProgress;
618        });
619    }
620
621    /// Finishes the progress bar and leaves the current message.
622    pub fn finish(&self) {
623        self.update_and_draw(|state| {
624            state.pos = state.len;
625            state.draw_next = state.pos;
626            state.status = Status::DoneVisible;
627        });
628    }
629
630    /// Finishes the progress bar at current position and leaves the current message.
631    pub fn finish_at_current_pos(&self) {
632        self.update_and_draw(|state| {
633            state.draw_next = state.pos;
634            state.status = Status::DoneVisible;
635        });
636    }
637
638    /// Finishes the progress bar and sets a message.
639    ///
640    /// For the message to be visible, `{msg}` placeholder
641    /// must be present in the template (see `ProgressStyle`).
642    pub fn finish_with_message(&self, msg: &str) {
643        let msg = msg.to_string();
644        self.update_and_draw(|state| {
645            state.message = msg;
646            state.pos = state.len;
647            state.draw_next = state.pos;
648            state.status = Status::DoneVisible;
649        });
650    }
651
652    /// Finishes the progress bar and completely clears it.
653    pub fn finish_and_clear(&self) {
654        self.update_and_draw(|state| {
655            state.pos = state.len;
656            state.draw_next = state.pos;
657            state.status = Status::DoneHidden;
658        });
659    }
660
661    /// Finishes the progress bar and leaves the current message and progress.
662    pub fn abandon(&self) {
663        self.update_and_draw(|state| {
664            state.status = Status::DoneVisible;
665        });
666    }
667
668    /// Finishes the progress bar and sets a message, and leaves the current progress.
669    ///
670    /// For the message to be visible, `{msg}` placeholder
671    /// must be present in the template (see `ProgressStyle`).
672    pub fn abandon_with_message(&self, msg: &str) {
673        let msg = msg.to_string();
674        self.update_and_draw(|state| {
675            state.message = msg;
676            state.status = Status::DoneVisible;
677        });
678    }
679
680    /// Sets a different draw target for the progress bar.
681    ///
682    /// This can be used to draw the progress bar to stderr
683    /// for instance:
684    ///
685    /// ```rust,no_run
686    /// # use imdl_indicatif::{ProgressBar, ProgressDrawTarget};
687    /// let pb = ProgressBar::new(100);
688    /// pb.set_draw_target(ProgressDrawTarget::stderr());
689    /// ```
690    pub fn set_draw_target(&self, target: ProgressDrawTarget) {
691        let mut state = self.state.write().unwrap();
692        state.draw_target.disconnect();
693        state.draw_target = target;
694    }
695
696    /// Wraps an iterator with the progress bar.
697    ///
698    /// ```rust,norun
699    /// # use imdl_indicatif::ProgressBar;
700    /// let v = vec![1, 2, 3];
701    /// let pb = ProgressBar::new(3);
702    /// for item in pb.wrap_iter(v.iter()) {
703    ///     // ...
704    /// }
705    /// ```
706    pub fn wrap_iter<It: Iterator>(&self, it: It) -> ProgressBarIter<It> {
707        ProgressBarIter {
708            bar: self.clone(),
709            it,
710        }
711    }
712
713    /// Wraps a Reader with the progress bar.
714    ///
715    /// ```rust,norun
716    /// # use std::fs::File;
717    /// # use std::io;
718    /// # use imdl_indicatif::ProgressBar;
719    /// # fn test () -> io::Result<()> {
720    /// let source = File::open("work.txt")?;
721    /// let mut target = File::create("done.txt")?;
722    /// let pb = ProgressBar::new(source.metadata()?.len());
723    /// io::copy(&mut pb.wrap_read(source), &mut target);
724    /// # Ok(())
725    /// # }
726    /// ```
727    pub fn wrap_read<R: io::Read>(&self, read: R) -> ProgressBarWrap<R> {
728        ProgressBarWrap {
729            bar: self.clone(),
730            wrap: read,
731        }
732    }
733
734    /// Wraps a Writer with the progress bar.
735    ///
736    /// ```rust,norun
737    /// # use std::fs::File;
738    /// # use std::io;
739    /// # use imdl_indicatif::ProgressBar;
740    /// # fn test () -> io::Result<()> {
741    /// let mut source = File::open("work.txt")?;
742    /// let target = File::create("done.txt")?;
743    /// let pb = ProgressBar::new(source.metadata()?.len());
744    /// io::copy(&mut source, &mut pb.wrap_write(target));
745    /// # Ok(())
746    /// # }
747    /// ```
748    pub fn wrap_write<W: io::Write>(&self, write: W) -> ProgressBarWrap<W> {
749        ProgressBarWrap {
750            bar: self.clone(),
751            wrap: write,
752        }
753    }
754
755    fn update_and_draw<F: FnOnce(&mut ProgressState)>(&self, f: F) {
756        let mut draw = false;
757        {
758            let mut state = self.state.write().unwrap();
759            let old_pos = state.pos;
760            f(&mut state);
761            let new_pos = state.pos;
762            if new_pos != old_pos {
763                state.est.record_step(new_pos);
764            }
765            if new_pos >= state.draw_next {
766                state.draw_next = new_pos.saturating_add(state.draw_delta);
767                draw = true;
768            }
769        }
770        if draw {
771            self.draw().ok();
772        }
773    }
774
775    fn draw(&self) -> io::Result<()> {
776        draw_state(&mut self.state.write().unwrap())
777    }
778
779    pub fn position(&self) -> u64 {
780        self.state.read().unwrap().pos
781    }
782}
783
784fn draw_state(state: &mut ProgressState) -> io::Result<()> {
785    // we can bail early if the draw target is hidden.
786    if state.draw_target.is_hidden() {
787        return Ok(());
788    }
789
790    let draw_state = ProgressDrawState {
791        lines: if state.should_render() {
792            state.style.format_state(&*state)
793        } else {
794            vec![]
795        },
796        orphan_lines: 0,
797        finished: state.is_finished(),
798        force_draw: false,
799        move_cursor: false,
800        ts: Instant::now(),
801    };
802    state.draw_target.apply_draw_state(draw_state)
803}
804
805impl Drop for ProgressState {
806    fn drop(&mut self) {
807        if self.is_finished() {
808            return;
809        }
810
811        self.status = Status::DoneHidden;
812        if self.pos >= self.draw_next {
813            self.draw_next = self.pos.saturating_add(self.draw_delta);
814            draw_state(self).ok();
815        }
816    }
817}
818
819#[test]
820fn test_pbar_zero() {
821    let pb = ProgressBar::new(0);
822    assert_eq!(pb.state.read().unwrap().fraction(), 1.0);
823}
824
825#[test]
826fn test_pbar_maxu64() {
827    let pb = ProgressBar::new(!0);
828    assert_eq!(pb.state.read().unwrap().fraction(), 0.0);
829}
830
831#[test]
832fn test_pbar_overflow() {
833    let pb = ProgressBar::new(1);
834    pb.set_draw_target(ProgressDrawTarget::hidden());
835    pb.inc(2);
836    pb.finish();
837}
838
839#[test]
840fn test_get_position() {
841    let pb = ProgressBar::new(1);
842    pb.set_draw_target(ProgressDrawTarget::hidden());
843    pb.inc(2);
844    let pos = pb.position();
845    assert_eq!(pos, 2);
846}
847
848struct MultiObject {
849    done: bool,
850    draw_state: Option<ProgressDrawState>,
851}
852
853struct MultiProgressState {
854    objects: Vec<MultiObject>,
855    draw_target: ProgressDrawTarget,
856    move_cursor: bool,
857}
858
859/// Manages multiple progress bars from different threads.
860pub struct MultiProgress {
861    state: RwLock<MultiProgressState>,
862    joining: AtomicBool,
863    tx: Sender<(usize, ProgressDrawState)>,
864    rx: Receiver<(usize, ProgressDrawState)>,
865}
866
867impl fmt::Debug for MultiProgress {
868    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
869        f.debug_struct("MultiProgress").finish()
870    }
871}
872
873unsafe impl Sync for MultiProgress {}
874
875impl Default for MultiProgress {
876    fn default() -> MultiProgress {
877        MultiProgress::with_draw_target(ProgressDrawTarget::stderr())
878    }
879}
880
881impl MultiProgress {
882    /// Creates a new multi progress object.
883    ///
884    /// Progress bars added to this object by default draw directly to stderr, and refresh
885    /// a maximum of 15 times a second. To change the refresh rate set the draw target to
886    /// one with a different refresh rate.
887    pub fn new() -> MultiProgress {
888        MultiProgress::default()
889    }
890
891    /// Creates a new multi progress object with the given draw target.
892    pub fn with_draw_target(draw_target: ProgressDrawTarget) -> MultiProgress {
893        let (tx, rx) = channel();
894        MultiProgress {
895            state: RwLock::new(MultiProgressState {
896                objects: vec![],
897                draw_target,
898                move_cursor: false,
899            }),
900            joining: AtomicBool::new(false),
901            tx,
902            rx,
903        }
904    }
905
906    /// Sets a different draw target for the multiprogress bar.
907    pub fn set_draw_target(&self, target: ProgressDrawTarget) {
908        let mut state = self.state.write().unwrap();
909        state.draw_target.disconnect();
910        state.draw_target = target;
911    }
912
913    /// Set whether we should try to move the cursor when possible instead of clearing lines.
914    ///
915    /// This can reduce flickering, but do not enable it if you intend to change the number of
916    /// progress bars.
917    pub fn set_move_cursor(&self, move_cursor: bool) {
918        self.state.write().unwrap().move_cursor = move_cursor;
919    }
920
921    /// Adds a progress bar.
922    ///
923    /// The progress bar added will have the draw target changed to a
924    /// remote draw target that is intercepted by the multi progress
925    /// object overriding custom `ProgressDrawTarget` settings.
926    pub fn add(&self, pb: ProgressBar) -> ProgressBar {
927        let mut state = self.state.write().unwrap();
928        let idx = state.objects.len();
929        state.objects.push(MultiObject {
930            done: false,
931            draw_state: None,
932        });
933        pb.set_draw_target(ProgressDrawTarget {
934            kind: ProgressDrawTargetKind::Remote(idx, Mutex::new(self.tx.clone())),
935        });
936        pb
937    }
938
939    /// Waits for all progress bars to report that they are finished.
940    ///
941    /// You need to call this as this will request the draw instructions
942    /// from the remote progress bars.  Not calling this will deadlock
943    /// your program.
944    pub fn join(&self) -> io::Result<()> {
945        self.join_impl(false)
946    }
947
948    /// Works like `join` but clears the progress bar in the end.
949    pub fn join_and_clear(&self) -> io::Result<()> {
950        self.join_impl(true)
951    }
952
953    fn is_done(&self) -> bool {
954        let state = self.state.read().unwrap();
955        if state.objects.is_empty() {
956            return true;
957        }
958        for obj in &state.objects {
959            if !obj.done {
960                return false;
961            }
962        }
963        true
964    }
965
966    fn join_impl(&self, clear: bool) -> io::Result<()> {
967        if self.joining.load(Ordering::Acquire) {
968            panic!("Already joining!");
969        }
970        self.joining.store(true, Ordering::Release);
971
972        let move_cursor = self.state.read().unwrap().move_cursor;
973        while !self.is_done() {
974            let (idx, draw_state) = self.rx.recv().unwrap();
975            let ts = draw_state.ts;
976            let force_draw = draw_state.finished || draw_state.force_draw;
977
978            let mut state = self.state.write().unwrap();
979            if draw_state.finished {
980                state.objects[idx].done = true;
981            }
982
983            // Split orphan lines out of the draw state, if any
984            let (orphan_lines, lines) = if draw_state.orphan_lines == 0 {
985                (vec![], draw_state.lines)
986            } else {
987                let split = draw_state.lines.split_at(draw_state.orphan_lines);
988                (split.0.to_vec(), split.1.to_vec())
989            };
990
991            let draw_state = ProgressDrawState {
992                lines,
993                orphan_lines: 0,
994                ..draw_state
995            };
996
997            state.objects[idx].draw_state = Some(draw_state);
998
999            // the rest from here is only drawing, we can skip it.
1000            if state.draw_target.is_hidden() {
1001                continue;
1002            }
1003
1004            let mut lines = vec![];
1005
1006            // Make orphaned lines appear at the top, so they can be properly
1007            // forgotten.
1008            let orphan_lines_count = orphan_lines.len();
1009            lines.extend(orphan_lines);
1010
1011            for obj in state.objects.iter() {
1012                if let Some(ref draw_state) = obj.draw_state {
1013                    lines.extend_from_slice(&draw_state.lines[..]);
1014                }
1015            }
1016
1017            // !any(!done) is also true when iter() is empty, contrary to all(done)
1018            let finished = !state.objects.iter().any(|ref x| !x.done);
1019            state.draw_target.apply_draw_state(ProgressDrawState {
1020                lines,
1021                orphan_lines: orphan_lines_count,
1022                force_draw,
1023                move_cursor,
1024                finished,
1025                ts,
1026            })?;
1027        }
1028
1029        if clear {
1030            let mut state = self.state.write().unwrap();
1031            state.draw_target.apply_draw_state(ProgressDrawState {
1032                lines: vec![],
1033                orphan_lines: 0,
1034                finished: true,
1035                force_draw: true,
1036                move_cursor,
1037                ts: Instant::now(),
1038            })?;
1039        }
1040
1041        self.joining.store(false, Ordering::Release);
1042
1043        Ok(())
1044    }
1045}
1046
1047/// Iterator for `wrap_iter`.
1048#[derive(Debug)]
1049pub struct ProgressBarIter<I> {
1050    bar: ProgressBar,
1051    it: I,
1052}
1053
1054impl<I: Iterator> Iterator for ProgressBarIter<I> {
1055    type Item = I::Item;
1056
1057    fn next(&mut self) -> Option<Self::Item> {
1058        let item = self.it.next();
1059
1060        if item.is_some() {
1061            self.bar.inc(1);
1062        }
1063
1064        item
1065    }
1066}
1067
1068/// wraps an io-object, either a Reader or a Writer (or both).
1069///
1070/// created by `wrap_read` or `wrap_write`
1071#[derive(Debug)]
1072pub struct ProgressBarWrap<W> {
1073    bar: ProgressBar,
1074    wrap: W,
1075}
1076
1077impl<R: io::Read> io::Read for ProgressBarWrap<R> {
1078    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
1079        let inc = self.wrap.read(buf)?;
1080        self.bar.inc(inc as u64);
1081        Ok(inc)
1082    }
1083}
1084
1085impl<S: io::Seek> io::Seek for ProgressBarWrap<S> {
1086    fn seek(&mut self, f: io::SeekFrom) -> io::Result<u64> {
1087        self.wrap.seek(f).map(|pos| {
1088            self.bar.set_position(pos);
1089            pos
1090        })
1091    }
1092}
1093
1094impl<W: io::Write> io::Write for ProgressBarWrap<W> {
1095    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
1096        self.wrap.write(buf).map(|inc| {
1097            self.bar.inc(inc as u64);
1098            inc
1099        })
1100    }
1101    fn flush(&mut self) -> io::Result<()> {
1102        self.wrap.flush()
1103    }
1104
1105    fn write_vectored(&mut self, bufs: &[io::IoSlice]) -> io::Result<usize> {
1106        self.wrap.write_vectored(bufs).map(|inc| {
1107            self.bar.inc(inc as u64);
1108            inc
1109        })
1110    }
1111    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
1112        self.wrap.write_all(buf).map(|()| {
1113            self.bar.inc(buf.len() as u64);
1114        })
1115    }
1116    // write_fmt can not be captured with reasonable effort.
1117    // as it uses write_all internally by default that should not be a problem.
1118    // fn write_fmt(&mut self, fmt: fmt::Arguments) -> io::Result<()>;
1119}
1120
1121#[cfg(test)]
1122mod tests {
1123    use super::*;
1124
1125    #[test]
1126    fn late_pb_drop() {
1127        let pb = ProgressBar::new(10);
1128        let mpb = MultiProgress::new();
1129        mpb.add(pb.clone());
1130    }
1131
1132    #[test]
1133    fn it_can_wrap_a_reader() {
1134        let bytes = &b"I am an implementation of io::Read"[..];
1135        let pb = ProgressBar::new(bytes.len() as u64);
1136        let mut reader = pb.wrap_read(bytes);
1137        let mut writer = Vec::new();
1138        io::copy(&mut reader, &mut writer).unwrap();
1139        assert_eq!(writer, bytes);
1140    }
1141
1142    #[test]
1143    fn it_can_wrap_a_writer() {
1144        let bytes = b"implementation of io::Read";
1145        let mut reader = &bytes[..];
1146        let pb = ProgressBar::new(bytes.len() as u64);
1147        let writer = Vec::new();
1148        let mut writer = pb.wrap_write(writer);
1149        io::copy(&mut reader, &mut writer).unwrap();
1150        assert_eq!(writer.wrap, bytes);
1151    }
1152
1153    #[test]
1154    fn progress_bar_sync_send() {
1155        let _: Box<dyn Sync> = Box::new(ProgressBar::new(1));
1156        let _: Box<dyn Send> = Box::new(ProgressBar::new(1));
1157        let _: Box<dyn Sync> = Box::new(MultiProgress::new());
1158        let _: Box<dyn Send> = Box::new(MultiProgress::new());
1159    }
1160}