hkxc/args/
progress_handler.rs

1use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
2use serde_hkx_features::progress::ProgressHandler;
3use std::path::Path;
4use std::sync::Arc;
5use std::sync::atomic::{AtomicU64, Ordering};
6
7/// For CLI progress bar
8#[derive(Debug, Clone)] // need clone for `tokio::spawn`
9pub struct CliProgressHandler {
10    progress_bar: ProgressBar,
11    success_bar: ProgressBar,
12    failure_bar: ProgressBar,
13    success_count: Arc<AtomicU64>,
14    failure_count: Arc<AtomicU64>,
15}
16
17impl CliProgressHandler {
18    /// Create a new `CliProgressHandler`.
19    pub fn new() -> Self {
20        let multi_progress = MultiProgress::new();
21        let progress_bar = multi_progress.add(ProgressBar::no_length());
22        let success_bar = multi_progress.add(ProgressBar::no_length());
23        let failure_bar = multi_progress.add(ProgressBar::no_length());
24
25        // Attempt to set the template and handle errors
26        match ProgressStyle::default_bar()
27            .template("{spinner:.green} [{elapsed_precise}] {bar:40.cyan/blue} {pos}/{len} {msg}")
28        {
29            Ok(style) => {
30                progress_bar.set_style(style.progress_chars("#>-"));
31            }
32            Err(e) => {
33                progress_bar.set_style(ProgressStyle::default_bar().progress_chars("#>-"));
34                tracing::error!("Failed to set progress bar template: {e}");
35            }
36        }
37
38        match ProgressStyle::default_spinner().template("✔ Success: {pos}") {
39            Ok(style) => success_bar.set_style(style),
40            Err(e) => {
41                success_bar.set_style(ProgressStyle::default_spinner());
42                tracing::error!("Failed to set success bar template: {e}");
43            }
44        }
45
46        match ProgressStyle::default_spinner().template("✖ Failed: {pos}") {
47            Ok(style) => failure_bar.set_style(style),
48            Err(e) => {
49                failure_bar.set_style(ProgressStyle::default_spinner());
50                tracing::error!("Failed to set failure bar template: {e}");
51            }
52        }
53
54        Self {
55            progress_bar,
56            success_bar,
57            failure_bar,
58            success_count: Arc::new(AtomicU64::new(0)),
59            failure_count: Arc::new(AtomicU64::new(0)),
60        }
61    }
62}
63
64impl ProgressHandler for CliProgressHandler {
65    fn on_empty(&self) {
66        println!(
67            "No files found in the directory to process. Please check if the directory contains valid files."
68        );
69    }
70
71    fn on_set_total(&mut self, total: usize) {
72        let total = total as u64;
73        self.progress_bar.set_length(total);
74        self.success_bar.set_length(total);
75        self.failure_bar.set_length(total);
76
77        let success_bar = self.success_bar.clone();
78        let failure_bar = self.failure_bar.clone();
79        let success_count = Arc::clone(&self.success_count);
80        let failure_count = Arc::clone(&self.failure_count);
81        std::thread::spawn({
82            move || loop {
83                success_bar.set_position(success_count.load(Ordering::Acquire));
84                failure_bar.set_position(failure_count.load(Ordering::Acquire));
85            }
86        });
87    }
88
89    fn inc(&self, progress: u64) {
90        self.progress_bar.inc(progress);
91    }
92
93    fn success_inc(&self, progress: u64) {
94        self.success_count.fetch_add(progress, Ordering::AcqRel);
95    }
96
97    fn failure_inc(&self, progress: u64) {
98        self.failure_count.fetch_add(progress, Ordering::AcqRel);
99    }
100
101    fn on_processing_path(&self, path: &Path) {
102        self.progress_bar.set_message(format!("{}", path.display()));
103    }
104
105    fn on_finish(&self) {
106        self.progress_bar.finish_with_message("All files processed");
107    }
108}