1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117

// this file contains a wrapper around pbr to work around three things:
//
// - there is no function to write above the progress bar
// - .draw() isn't exposed so we can't bypass the ratelimit in tick.
//   This means we can't reliably redraw the graph after we wrote above it.
//   We have to implement rate limiting in our wrapper to ensure we are
//   able to bypass it when needed.
// - using colored strings breaks pbr
//
// https://github.com/a8m/pb/pull/62

use pbr;
use atty;
use colored::Colorize;
use std::fmt::Display;
use std::io::prelude::*;
use std::io::{self, Stdout};
use time::{self, SteadyTime, Duration};


macro_rules! printfl {
   ($w:expr, $($tt:tt)*) => {{
        $w.write(&format!($($tt)*).as_bytes()).ok().expect("write() fail");
        $w.flush().ok().expect("flush() fail");
    }}
}

pub struct ProgressBar {
    pb: pbr::ProgressBar<Stdout>,
    current: u64,
    last_refresh_time: SteadyTime,
    max_refresh_rate: Option<time::Duration>,
    atty: bool,
}

impl ProgressBar {
    #[inline]
    pub fn new(total: u64) -> ProgressBar {
        let mut pb = pbr::ProgressBar::new(total);
        pb.format("(=> )");

        let now = SteadyTime::now();
        let refresh_rate = Duration::milliseconds(250);
        let atty = atty::is(atty::Stream::Stdout);

        ProgressBar {
            pb,
            current: 0,
            last_refresh_time: now - refresh_rate,
            max_refresh_rate: Some(refresh_rate),
            atty,
        }
    }

    #[inline]
    pub fn draw(&mut self) {
        if !self.atty {
            return;
        }

        self.pb.tick()
    }

    #[inline]
    pub fn print_help(&mut self) {
        self.writeln(format!("{} {}", "[+]".bold(),
            "[h] help, [p] pause, [r] resume, [+] increase threads, [-] decrease threads".dimmed()));
    }

    #[inline]
    pub fn writeln<T: Display>(&mut self, s: T) {
        printfl!(io::stderr(), "\r\x1B[2K{}\n", s);
        self.draw()
    }

    #[inline]
    pub fn tick(&mut self) {
        let now = SteadyTime::now();
        if let Some(mrr) = self.max_refresh_rate {
            if now - self.last_refresh_time < mrr {
                return;
            }
        }

        self.draw();

        self.last_refresh_time = SteadyTime::now();
    }

    #[inline]
    pub fn inc(&mut self) {
        if !self.atty {
            return;
        }

        let now = SteadyTime::now();
        if let Some(mrr) = self.max_refresh_rate {
            if now - self.last_refresh_time < mrr {
                self.current += 1;
                return;
            }
        }

        self.pb.set(self.current);

        self.last_refresh_time = SteadyTime::now();
    }

    #[inline]
    pub fn finish_replace<T: Display>(&self, s: T) {
        if self.atty {
            print!("\r\x1B[2K{}", s);
        } else {
            print!("{}", s);
        }
    }
}