1use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
2use indicatif_log_bridge::LogWrapper;
3use std::io::{IsTerminal, stderr};
4use std::path::Path;
5use std::time::Duration;
6
7const SPINNER_STYLE: &str = "{elapsed_precise} {spinner:.green} {pos:>7} {msg}";
8const NORMAL_STYLE: &str =
9 "{elapsed_precise} +{eta:>3} {percent:>3}% {bar:40.cyan/blue} {pos:>7}/{len:7} {msg}";
10const FILE_STYLE: &str = " {elapsed:>3} +{eta:>3} {percent:>3}% {bar:10.cyan/blue} {wide_msg}";
11
12#[derive(Debug)]
13pub(crate) struct Progress {
14 inner: Option<ProgressBar>,
15 multi: Option<MultiProgress>,
16}
17
18impl Progress {
19 pub fn none() -> Self {
20 Self {
21 inner: None,
22 multi: None,
23 }
24 }
25
26 pub fn set_message(&self, msg: impl Into<String>) {
27 if let Some(inner) = &self.inner {
28 inner.set_message(msg.into());
29 }
30 }
31
32 pub fn inc(&self, amount: u64) {
33 if let Some(inner) = &self.inner {
34 inner.inc(amount);
35 }
36 }
37
38 pub fn length(&self) -> Option<u64> {
39 self.inner.as_ref().and_then(|inner| inner.length())
40 }
41
42 pub fn set_length(&self, len: u64) {
43 if let Some(inner) = &self.inner {
44 if inner.length().is_none() {
45 inner.set_style(ProgressStyle::with_template(NORMAL_STYLE).unwrap());
46 }
47 inner.set_length(len);
48 }
49 }
50
51 pub fn finish(&self) {
52 if let Some(inner) = &self.inner {
53 inner.finish();
54 if let Some(multi) = &self.multi {
55 multi.remove(inner);
56 }
57 }
58 }
59
60 pub fn suspend_for<F, R, S: IsTerminal>(&self, stream: S, f: F) -> R
61 where
62 F: FnOnce() -> R,
63 {
64 if !stream.is_terminal() {
65 return f();
66 }
67 if let Some(multi) = &self.multi {
68 multi.suspend(f)
69 } else if let Some(inner) = &self.inner {
70 inner.suspend(f)
71 } else {
72 f()
73 }
74 }
75}
76
77#[derive(Debug)]
78pub struct ProgressBuilder {
79 multi: MultiProgress,
80 pub is_enabled: bool,
81 pub is_file_enabled: bool,
82}
83
84impl Default for ProgressBuilder {
85 fn default() -> Self {
86 Self {
87 multi: MultiProgress::default(),
88 is_enabled: stderr().is_terminal(),
89 is_file_enabled: false,
90 }
91 }
92}
93
94impl ProgressBuilder {
95 pub fn new() -> Self {
96 Self::default()
97 }
98
99 pub fn init_logger(&self, logger: env_logger::Logger) -> anyhow::Result<()> {
100 let max_level = logger.filter();
101 LogWrapper::new(self.multi.clone(), logger).try_init()?;
102 log::set_max_level(max_level);
103 Ok(())
104 }
105
106 pub(crate) fn add_spinner(&self) -> Progress {
107 if !self.is_enabled {
108 return Progress::none();
109 }
110 let progress = self.multi.add(ProgressBar::new_spinner());
111 progress.enable_steady_tick(Duration::from_secs(1));
112 progress.set_style(ProgressStyle::with_template(SPINNER_STYLE).unwrap());
113 Progress {
114 inner: Some(progress),
115 multi: Some(self.multi.clone()),
116 }
117 }
118
119 pub(crate) fn add_file(&self, path: &Path, file_size: u64) -> Progress {
120 if !self.is_enabled || !self.is_file_enabled {
121 return Progress::none();
122 }
123 let progress = self.multi.add(ProgressBar::new(file_size));
124 progress.set_style(ProgressStyle::with_template(FILE_STYLE).unwrap());
125 if let Some(parent) = path.parent()
126 && let Some(file_name) = path.file_name()
127 {
128 progress.set_message(format!(
129 "{} ({})",
130 file_name.to_string_lossy(),
131 parent.to_string_lossy()
132 ));
133 } else {
134 progress.set_message(path.to_string_lossy().to_string());
135 }
136 Progress {
137 inner: Some(progress),
138 multi: Some(self.multi.clone()),
139 }
140 }
141}