nils_term/
progress.rs

1use std::io::{self, IsTerminal};
2use std::sync::atomic::{AtomicBool, Ordering};
3use std::sync::{Arc, Mutex};
4
5use indicatif::{ProgressBar, ProgressStyle};
6
7#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8pub enum ProgressEnabled {
9    /// Enable progress only when the selected draw target is a TTY.
10    ///
11    /// Notes:
12    /// - With the default draw target (stderr), this disables progress when stderr is not a TTY,
13    ///   keeping stdout clean for piping and machine-readable output.
14    /// - With `ProgressDrawTarget::to_writer(...)` (tests), auto-enable is never blocked by TTY
15    ///   detection.
16    Auto,
17    On,
18    Off,
19}
20
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub enum ProgressFinish {
23    /// Leave the progress line visible when finished.
24    Leave,
25    /// Clear the progress line when finished.
26    Clear,
27}
28
29#[derive(Debug, Clone)]
30pub struct ProgressOptions {
31    pub enabled: ProgressEnabled,
32    pub prefix: String,
33    /// Fixed terminal width used by writer-based draw targets (tests).
34    pub width: Option<u16>,
35    pub finish: ProgressFinish,
36    pub draw_target: ProgressDrawTarget,
37}
38
39impl Default for ProgressOptions {
40    fn default() -> Self {
41        Self {
42            enabled: ProgressEnabled::Auto,
43            prefix: String::new(),
44            width: None,
45            finish: ProgressFinish::Leave,
46            draw_target: ProgressDrawTarget::stderr(),
47        }
48    }
49}
50
51impl ProgressOptions {
52    pub fn with_enabled(mut self, enabled: ProgressEnabled) -> Self {
53        self.enabled = enabled;
54        self
55    }
56
57    pub fn with_prefix(mut self, prefix: impl Into<String>) -> Self {
58        self.prefix = prefix.into();
59        self
60    }
61
62    pub fn with_width(mut self, width: Option<u16>) -> Self {
63        self.width = width;
64        self
65    }
66
67    pub fn with_finish(mut self, finish: ProgressFinish) -> Self {
68        self.finish = finish;
69        self
70    }
71
72    pub fn with_draw_target(mut self, draw_target: ProgressDrawTarget) -> Self {
73        self.draw_target = draw_target;
74        self
75    }
76}
77
78#[derive(Debug, Clone)]
79pub enum ProgressDrawTarget {
80    /// Draw to stderr (default).
81    Stderr,
82    /// Draw to an in-memory writer (intended for deterministic tests).
83    Writer { buffer: Arc<Mutex<Vec<u8>>> },
84}
85
86impl ProgressDrawTarget {
87    pub fn stderr() -> Self {
88        Self::Stderr
89    }
90
91    pub fn to_writer(buffer: Arc<Mutex<Vec<u8>>>) -> Self {
92        Self::Writer { buffer }
93    }
94}
95
96#[derive(Debug, Clone)]
97pub struct Progress {
98    state: Option<Arc<ProgressState>>,
99}
100
101#[derive(Debug)]
102struct ProgressState {
103    bar: ProgressBar,
104    finish: ProgressFinish,
105    rendered: AtomicBool,
106    finished: AtomicBool,
107}
108
109impl Drop for ProgressState {
110    fn drop(&mut self) {
111        let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
112            if self.finished.load(Ordering::Relaxed) {
113                return;
114            }
115            if !self.rendered.load(Ordering::Relaxed) {
116                return;
117            }
118            match self.finish {
119                ProgressFinish::Leave => self.bar.finish(),
120                ProgressFinish::Clear => self.bar.finish_and_clear(),
121            }
122        }));
123    }
124}
125
126impl Progress {
127    /// Create a determinate progress bar.
128    pub fn new(total: u64, options: ProgressOptions) -> Self {
129        if !should_enable(&options) {
130            return Self { state: None };
131        }
132
133        let draw_target = to_indicatif_draw_target(&options.draw_target, options.width);
134
135        let bar = ProgressBar::new(total);
136        bar.set_draw_target(draw_target);
137        bar.set_style(determinate_style());
138
139        let state = Arc::new(ProgressState {
140            bar,
141            finish: options.finish,
142            rendered: AtomicBool::new(false),
143            finished: AtomicBool::new(false),
144        });
145
146        if !options.prefix.is_empty() {
147            state.rendered.store(true, Ordering::Relaxed);
148            state.bar.set_prefix(options.prefix);
149        }
150
151        Self { state: Some(state) }
152    }
153
154    /// Create a spinner progress indicator.
155    pub fn spinner(options: ProgressOptions) -> Self {
156        if !should_enable(&options) {
157            return Self { state: None };
158        }
159
160        let draw_target = to_indicatif_draw_target(&options.draw_target, options.width);
161
162        let bar = ProgressBar::new_spinner();
163        bar.set_draw_target(draw_target);
164        bar.set_style(spinner_style());
165
166        let state = Arc::new(ProgressState {
167            bar,
168            finish: options.finish,
169            rendered: AtomicBool::new(false),
170            finished: AtomicBool::new(false),
171        });
172
173        if !options.prefix.is_empty() {
174            state.rendered.store(true, Ordering::Relaxed);
175            state.bar.set_prefix(options.prefix);
176        }
177
178        Self { state: Some(state) }
179    }
180
181    pub fn set_position(&self, pos: u64) {
182        if let Some(state) = &self.state {
183            state.rendered.store(true, Ordering::Relaxed);
184            state.bar.set_position(pos);
185        }
186    }
187
188    pub fn inc(&self, delta: u64) {
189        if let Some(state) = &self.state {
190            state.rendered.store(true, Ordering::Relaxed);
191            state.bar.inc(delta);
192        }
193    }
194
195    pub fn tick(&self) {
196        if let Some(state) = &self.state {
197            state.rendered.store(true, Ordering::Relaxed);
198            state.bar.tick();
199        }
200    }
201
202    pub fn set_message(&self, message: impl Into<String>) {
203        if let Some(state) = &self.state {
204            state.rendered.store(true, Ordering::Relaxed);
205            state.bar.set_message(message.into());
206        }
207    }
208
209    /// Finish and leave the progress output visible.
210    pub fn finish(&self) {
211        if let Some(state) = &self.state {
212            if state.finished.swap(true, Ordering::Relaxed) {
213                return;
214            }
215            state.rendered.store(true, Ordering::Relaxed);
216            state.bar.finish();
217        }
218    }
219
220    pub fn finish_with_message(&self, message: impl Into<String>) {
221        if let Some(state) = &self.state {
222            if state.finished.swap(true, Ordering::Relaxed) {
223                return;
224            }
225            state.rendered.store(true, Ordering::Relaxed);
226            state.bar.finish_with_message(message.into());
227        }
228    }
229
230    /// Finish and clear the progress output.
231    pub fn finish_and_clear(&self) {
232        if let Some(state) = &self.state {
233            if state.finished.swap(true, Ordering::Relaxed) {
234                return;
235            }
236            state.rendered.store(true, Ordering::Relaxed);
237            state.bar.finish_and_clear();
238        }
239    }
240
241    pub fn suspend<F: FnOnce() -> R, R>(&self, f: F) -> R {
242        match &self.state {
243            Some(state) => state.bar.suspend(f),
244            None => f(),
245        }
246    }
247}
248
249fn should_enable(options: &ProgressOptions) -> bool {
250    match options.enabled {
251        ProgressEnabled::On => true,
252        ProgressEnabled::Off => false,
253        ProgressEnabled::Auto => match &options.draw_target {
254            ProgressDrawTarget::Stderr => io::stderr().is_terminal(),
255            ProgressDrawTarget::Writer { .. } => true,
256        },
257    }
258}
259
260fn determinate_style() -> ProgressStyle {
261    // Keep this stable and non-panicking even if indicatif changes behavior.
262    let style = ProgressStyle::with_template("{prefix}{wide_bar} {pos}/{len} {msg}");
263    match style {
264        Ok(style) => style.progress_chars("#-"),
265        Err(_) => ProgressStyle::default_bar().progress_chars("#-"),
266    }
267}
268
269fn spinner_style() -> ProgressStyle {
270    let style = ProgressStyle::with_template("{prefix}{spinner} {msg}");
271    match style {
272        Ok(style) => style.tick_chars(r"-\|/"),
273        Err(_) => ProgressStyle::default_spinner().tick_chars(r"-\|/"),
274    }
275}
276
277fn to_indicatif_draw_target(
278    draw_target: &ProgressDrawTarget,
279    width: Option<u16>,
280) -> indicatif::ProgressDrawTarget {
281    match draw_target {
282        ProgressDrawTarget::Stderr => indicatif::ProgressDrawTarget::stderr(),
283        ProgressDrawTarget::Writer { buffer } => indicatif::ProgressDrawTarget::term_like(
284            Box::new(WriterTerm::new(buffer.clone(), width.unwrap_or(80))),
285        ),
286    }
287}
288
289#[derive(Debug)]
290struct WriterTerm {
291    buffer: Arc<Mutex<Vec<u8>>>,
292    width: u16,
293}
294
295impl WriterTerm {
296    fn new(buffer: Arc<Mutex<Vec<u8>>>, width: u16) -> Self {
297        Self { buffer, width }
298    }
299
300    fn write_all(&self, bytes: &[u8]) -> io::Result<()> {
301        let mut guard = match self.buffer.lock() {
302            Ok(guard) => guard,
303            Err(poisoned) => poisoned.into_inner(),
304        };
305        guard.extend_from_slice(bytes);
306        Ok(())
307    }
308
309    fn write_str_bytes(&self, s: &str) -> io::Result<()> {
310        self.write_all(s.as_bytes())
311    }
312}
313
314impl indicatif::TermLike for WriterTerm {
315    fn width(&self) -> u16 {
316        self.width
317    }
318
319    fn move_cursor_up(&self, n: usize) -> io::Result<()> {
320        self.write_str_bytes(&format!("\u{1b}[{n}A"))
321    }
322
323    fn move_cursor_down(&self, n: usize) -> io::Result<()> {
324        self.write_str_bytes(&format!("\u{1b}[{n}B"))
325    }
326
327    fn move_cursor_right(&self, n: usize) -> io::Result<()> {
328        self.write_str_bytes(&format!("\u{1b}[{n}C"))
329    }
330
331    fn move_cursor_left(&self, n: usize) -> io::Result<()> {
332        self.write_str_bytes(&format!("\u{1b}[{n}D"))
333    }
334
335    fn write_line(&self, s: &str) -> io::Result<()> {
336        self.write_str_bytes(s)?;
337        self.write_all(b"\n")
338    }
339
340    fn write_str(&self, s: &str) -> io::Result<()> {
341        self.write_str_bytes(s)
342    }
343
344    fn clear_line(&self) -> io::Result<()> {
345        self.write_str_bytes("\r\u{1b}[2K")
346    }
347
348    fn flush(&self) -> io::Result<()> {
349        Ok(())
350    }
351}
nils_term/progress.rs

nils_term/
progress.rs
