cached_path/
progress_bar.rs

1use std::io::{self, Write};
2use std::time::Instant;
3
4/// Progress bar types.
5///
6/// This can be set with
7/// [`CacheBuilder::progress_bar()`](struct.CacheBuilder.html#method.progress_bar).
8#[derive(Debug, Clone, Default)]
9pub enum ProgressBar {
10    /// Gives pretty, verbose progress bars.
11    #[default]
12    Full,
13    /// Gives progress bars with minimal output.
14    ///
15    /// This is a good option to use if your output is being captured to a file but you
16    /// still want to see progress updates.
17    Light,
18}
19
20impl ProgressBar {
21    pub(crate) fn wrap_download<W: Write>(
22        &self,
23        resource: &str,
24        content_length: Option<u64>,
25        writer: W,
26    ) -> DownloadWrapper<W> {
27        let bar: Box<dyn DownloadBar> = match self {
28            ProgressBar::Full => Box::new(FullDownloadBar::new(content_length)),
29            ProgressBar::Light => Box::new(LightDownloadBar::new(resource, content_length)),
30        };
31        DownloadWrapper::new(bar, writer)
32    }
33}
34
35pub(crate) struct DownloadWrapper<W: Write> {
36    bar: Box<dyn DownloadBar>,
37    writer: W,
38}
39
40impl<W> DownloadWrapper<W>
41where
42    W: Write,
43{
44    fn new(bar: Box<dyn DownloadBar>, writer: W) -> Self {
45        Self { bar, writer }
46    }
47
48    pub(crate) fn finish(&self) {
49        self.bar.finish();
50    }
51}
52
53impl<W: Write> Write for DownloadWrapper<W> {
54    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
55        self.writer.write(buf)
56    }
57
58    fn flush(&mut self) -> io::Result<()> {
59        self.writer.flush()
60    }
61
62    fn write_vectored(&mut self, bufs: &[io::IoSlice]) -> io::Result<usize> {
63        self.writer.write_vectored(bufs)
64    }
65
66    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
67        self.writer.write_all(buf).map(|()| {
68            self.bar.tick(buf.len());
69        })
70    }
71}
72
73trait DownloadBar {
74    fn tick(&mut self, chunk_size: usize);
75
76    fn finish(&self);
77}
78
79pub(crate) struct FullDownloadBar {
80    bar: indicatif::ProgressBar,
81}
82
83impl FullDownloadBar {
84    pub(crate) fn new(content_length: Option<u64>) -> Self {
85        let bar = match content_length {
86            Some(length) => {
87                let bar = indicatif::ProgressBar::new(length);
88                bar.set_style(
89                    indicatif::ProgressStyle::default_bar()
90                    .progress_chars("=>-")
91                    .template(
92                        "{msg:.bold.cyan/blue} [{bar:20.cyan/blue}][{percent}%] {bytes}/{total_bytes:.bold} |{bytes_per_sec}|",
93                    )
94                    .expect("Failed to set progress bar template")
95                );
96                bar
97            }
98            None => {
99                let bar = indicatif::ProgressBar::new_spinner();
100                bar.set_style(
101                    indicatif::ProgressStyle::default_bar()
102                        .tick_strings(&[
103                            "⠁⠁⠉⠙⠚⠒⠂⠂⠒⠲⠴⠤⠄⠄⠤⠠⠠⠤⠦⠖",
104                            "⠉⠙⠚⠒⠂⠂⠒⠲⠴⠤⠄⠄⠤⠠⠠⠤⠦⠖⠒⠐",
105                            "⠚⠒⠂⠂⠒⠲⠴⠤⠄⠄⠤⠠⠠⠤⠦⠖⠒⠐⠐⠒",
106                            "⠂⠂⠒⠲⠴⠤⠄⠄⠤⠠⠠⠤⠦⠖⠒⠐⠐⠒⠓⠋",
107                            "⠒⠲⠴⠤⠄⠄⠤⠠⠠⠤⠦⠖⠒⠐⠐⠒⠓⠋⠉⠈",
108                            "⠴⠤⠄⠄⠤⠠⠠⠤⠦⠖⠒⠐⠐⠒⠓⠋⠉⠈⠈⠉",
109                            "⠄⠄⠤⠠⠠⠤⠦⠖⠒⠐⠐⠒⠓⠋⠉⠈⠈⠉⠙⠚",
110                            "⠤⠠⠠⠤⠦⠖⠒⠐⠐⠒⠓⠋⠉⠈⠈⠉⠙⠚⠒⠂",
111                            "⠠⠤⠦⠖⠒⠐⠐⠒⠓⠋⠉⠈⠈⠉⠙⠚⠒⠂⠂⠒",
112                            "⠦⠖⠒⠐⠐⠒⠓⠋⠉⠈⠈⠉⠙⠚⠒⠂⠂⠒⠲⠴",
113                            "⠒⠐⠐⠒⠓⠋⠉⠈⠈⠉⠙⠚⠒⠂⠂⠒⠲⠴⠤⠄",
114                            "⠐⠒⠓⠋⠉⠈⠈⠉⠙⠚⠒⠂⠂⠒⠲⠴⠤⠄⠄⠤",
115                            "⠓⠋⠉⠈⠈⠉⠙⠚⠒⠂⠂⠒⠲⠴⠤⠄⠄⠤⠠⠠",
116                        ])
117                        .template(
118                            "{msg:.bold.cyan/blue} [{spinner:.cyan/blue}] {bytes:.bold} |{bytes_per_sec}|",
119                        )
120                        .expect("Failed to set spinner progress bar template"),
121                );
122                bar
123            }
124        };
125        bar.set_message("Downloading");
126        // Update every 1 MBs.
127        // NOTE: If we don't set this, the updates happen WAY too frequently and it makes downloads
128        // take about twice as long.
129        // In indicatif 0.17, set_draw_delta was removed. We can use enable_steady_tick instead
130        // or rely on the default drawing behavior which should be more efficient.
131        bar.enable_steady_tick(std::time::Duration::from_millis(100));
132        Self { bar }
133    }
134}
135
136impl DownloadBar for FullDownloadBar {
137    fn tick(&mut self, chunk_size: usize) {
138        self.bar.inc(chunk_size as u64);
139    }
140
141    fn finish(&self) {
142        self.bar.set_message("Downloaded");
143        self.bar.set_style(
144            indicatif::ProgressStyle::default_bar()
145                .template("{msg:.green.bold} {total_bytes:.bold} in {elapsed}")
146                .expect("Failed to set finish progress bar template"),
147        );
148        self.bar.finish();
149    }
150}
151
152pub(crate) struct LightDownloadBar {
153    start_time: Instant,
154    bytes: usize,
155    bytes_since_last_update: usize,
156}
157
158impl LightDownloadBar {
159    pub(crate) fn new(resource: &str, content_length: Option<u64>) -> Self {
160        if let Some(size) = content_length {
161            eprint!(
162                "Downloading {resource} [{}]...",
163                indicatif::HumanBytes(size)
164            );
165        } else {
166            eprint!("Downloading {resource}...");
167        }
168        io::stderr().flush().ok();
169        Self {
170            start_time: Instant::now(),
171            bytes: 0,
172            bytes_since_last_update: 0,
173        }
174    }
175}
176
177impl DownloadBar for LightDownloadBar {
178    fn tick(&mut self, chunk_size: usize) {
179        self.bytes_since_last_update += chunk_size;
180        // Update every 100 MBs.
181        if self.bytes_since_last_update > 100_000_000 {
182            eprint!(".");
183            io::stderr().flush().ok();
184            self.bytes_since_last_update = 0;
185        }
186        self.bytes += chunk_size;
187    }
188
189    fn finish(&self) {
190        let duration = Instant::now().duration_since(self.start_time);
191        eprintln!(
192            " ✓ Done! Finished in {}",
193            indicatif::HumanDuration(duration)
194        );
195        io::stderr().flush().ok();
196    }
197}