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_STYLE0: &str = "{elapsed_precise} {spinner:.green} ";
8const SPINNER_STYLE1_NUM: &str = "{pos:>7} {msg}";
9const SPINNER_STYLE1_SIZE: &str = "{bytes:>7} {msg}";
10const NORMAL_STYLE0: &str = "{elapsed_precise} +{eta:>3} {percent:>3}% {bar:40.cyan/blue} ";
11const NORMAL_STYLE1_NUM: &str = "{pos:>7}/{len:7} {msg}";
12const NORMAL_STYLE1_SIZE: &str = "{bytes:>7}/{total_bytes:7} {msg}";
13const FILE_STYLE: &str = " {elapsed:>3} +{eta:>3} {percent:>3}% {bar:10.cyan/blue} {wide_msg}";
14
15#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
16pub(crate) struct ProgressValue {
17 pub(crate) num_files: u64,
18 size: u64,
19}
20
21impl ProgressValue {
22 pub(crate) fn with_size(size: u64) -> Self {
23 Self { size, num_files: 1 }
24 }
25
26 pub(crate) fn with_skip(size: u64) -> Self {
27 Self { size, num_files: 1 }
28 }
29}
30
31impl std::ops::Add for ProgressValue {
32 type Output = Self;
33
34 fn add(self, other: Self) -> Self {
35 Self {
36 num_files: self.num_files + other.num_files,
37 size: self.size + other.size,
38 }
39 }
40}
41
42impl std::ops::AddAssign for ProgressValue {
43 fn add_assign(&mut self, other: Self) {
44 self.num_files += other.num_files;
45 self.size += other.size;
46 }
47}
48
49#[derive(Debug, Default)]
50pub(crate) struct Progress {
51 inner: Option<ProgressBar>,
52 pos: ProgressValue,
53 use_bytes: bool,
54 len: Option<ProgressValue>,
55 multi: Option<MultiProgress>,
56}
57
58impl Progress {
59 pub fn none() -> Self {
60 Self {
61 inner: None,
62 multi: None,
63 ..Default::default()
64 }
65 }
66
67 fn update_style(&self) {
68 if let Some(inner) = &self.inner {
69 let style = if self.len.is_some() {
70 if self.use_bytes {
71 format!("{NORMAL_STYLE0}{NORMAL_STYLE1_SIZE}")
72 } else {
73 format!("{NORMAL_STYLE0}{NORMAL_STYLE1_NUM}")
74 }
75 } else {
76 if self.use_bytes {
77 format!("{SPINNER_STYLE0}{SPINNER_STYLE1_SIZE}")
78 } else {
79 format!("{SPINNER_STYLE0}{SPINNER_STYLE1_NUM}")
80 }
81 };
82 inner.set_style(ProgressStyle::with_template(&style).unwrap());
83 }
84 }
85
86 fn update_position(&self) {
87 if let Some(inner) = &self.inner {
88 inner.set_position(if self.use_bytes {
89 self.pos.size
90 } else {
91 self.pos.num_files
92 });
93 }
94 }
95
96 fn update_length(&self) {
97 if let Some(inner) = &self.inner
98 && let Some(len) = self.len
99 {
100 inner.set_length(if self.use_bytes {
101 len.size
102 } else {
103 len.num_files
104 });
105 }
106 }
107
108 pub fn set_message(&self, msg: impl Into<String>) {
109 if let Some(inner) = &self.inner {
110 inner.set_message(msg.into());
111 }
112 }
113
114 pub fn inc(&mut self, amount: ProgressValue) {
115 self.pos += amount;
116 self.update_position();
117 }
118
119 pub fn set_length(&mut self, len: ProgressValue) {
120 self.len = Some(len);
121 self.update_style();
122 self.update_length();
123 }
124
125 pub fn use_bytes(&mut self) {
126 self.use_bytes = true;
127 self.update_style();
128 self.update_length();
129 self.update_position();
130 }
131
132 pub fn finish(&self) {
133 if let Some(inner) = &self.inner {
134 inner.finish();
135 if let Some(multi) = &self.multi {
136 multi.remove(inner);
137 }
138 }
139 }
140
141 pub fn suspend_for<F, R, S: IsTerminal>(&self, stream: S, f: F) -> R
142 where
143 F: FnOnce() -> R,
144 {
145 if !stream.is_terminal() {
146 return f();
147 }
148 if let Some(multi) = &self.multi {
149 multi.suspend(f)
150 } else if let Some(inner) = &self.inner {
151 inner.suspend(f)
152 } else {
153 f()
154 }
155 }
156}
157
158#[derive(Debug)]
159pub struct ProgressBuilder {
160 multi: MultiProgress,
161 pub is_enabled: bool,
162 pub is_file_enabled: bool,
163}
164
165impl Default for ProgressBuilder {
166 fn default() -> Self {
167 Self {
168 multi: MultiProgress::default(),
169 is_enabled: stderr().is_terminal(),
170 is_file_enabled: false,
171 }
172 }
173}
174
175impl ProgressBuilder {
176 pub fn new() -> Self {
177 Self::default()
178 }
179
180 pub fn init_logger(&self, logger: env_logger::Logger) -> anyhow::Result<()> {
181 let max_level = logger.filter();
182 LogWrapper::new(self.multi.clone(), logger).try_init()?;
183 log::set_max_level(max_level);
184 Ok(())
185 }
186
187 pub(crate) fn add_spinner(&self) -> Progress {
188 if !self.is_enabled {
189 return Progress::none();
190 }
191 let inner = self.multi.add(ProgressBar::new_spinner());
192 inner.enable_steady_tick(Duration::from_secs(1));
193 let progress = Progress {
194 inner: Some(inner),
195 multi: Some(self.multi.clone()),
196 ..Default::default()
197 };
198 progress.update_style();
199 progress
200 }
201
202 pub(crate) fn add_file(&self, path: &Path, file_size: u64) -> Progress {
203 if !self.is_enabled || !self.is_file_enabled {
204 return Progress::none();
205 }
206 let inner = self.multi.add(ProgressBar::new(file_size));
207 inner.set_style(ProgressStyle::with_template(FILE_STYLE).unwrap());
208 if let Some(parent) = path.parent()
209 && let Some(file_name) = path.file_name()
210 {
211 inner.set_message(format!(
212 "{} ({})",
213 file_name.to_string_lossy(),
214 parent.to_string_lossy()
215 ));
216 } else {
217 inner.set_message(path.to_string_lossy().to_string());
218 }
219 Progress {
220 inner: Some(inner),
221 multi: Some(self.multi.clone()),
222 use_bytes: true,
223 ..Default::default()
224 }
225 }
226}