ibdl_core/progress_bars/
mod.rs

1use ibdl_common::{
2    tokio::{spawn, sync::mpsc::Receiver},
3    ImageBoards,
4};
5use indicatif::{
6    HumanBytes, MultiProgress, ProgressBar, ProgressDrawTarget, ProgressState, ProgressStyle,
7};
8use std::{
9    fmt::Write,
10    sync::{
11        atomic::{AtomicU64, AtomicUsize, Ordering},
12        Arc,
13    },
14    time::Duration,
15};
16
17const PROGRESS_CHARS: &str = "━━";
18
19struct BarTemplates {
20    pub main: &'static str,
21    pub download: &'static str,
22}
23
24impl BarTemplates {
25    /// Returns special-themed progress bar templates for each variant
26    #[inline]
27    pub fn new(imageboard: ImageBoards) -> Self {
28        match imageboard {
29            ImageBoards::E621 => Self {
30                main: "{spinner:.yellow.bold} {elapsed_precise:.bold} {wide_bar:.blue/white.dim} {percent:.bold}  {pos:.yellow} (eta. {eta})",
31                download: "{spinner:.blue.bold} {bar:40.yellow/white.dim} {percent:.bold} | {byte_progress:21.blue} @ {bytes_per_sec:>13.yellow} (eta. {eta:<4.blue})",
32            },
33            ImageBoards::GelbooruV0_2 => Self {
34                main: "{spinner:.red.bold} {elapsed_precise:.bold} {wide_bar:.red/white.dim} {percent:.bold}  {pos:.bold} (eta. {eta})", 
35                download: "{spinner:.red.bold} {bar:40.red/white.dim} {percent:.bold} | {byte_progress:21.bold.green} @ {bytes_per_sec:>13.red} (eta. {eta:<4})",
36            },
37            _ => Self::default(),
38        }
39    }
40}
41
42impl Default for BarTemplates {
43    fn default() -> Self {
44        Self {
45            main: "{spinner:.green.bold} {elapsed_precise:.bold} {wide_bar:.green/white.dim} {percent:.bold}  {pos:.green} (eta. {eta:.blue})",
46            download: "{spinner:.green.bold} {bar:40.green/white.dim} {percent:.bold} | {byte_progress:21.green} @ {bytes_per_sec:>13.red} (eta. {eta:<4.blue})",
47        }
48    }
49}
50
51/// Struct to condense a commonly used duo of progress bar instances and counters for downloaded posts.
52///
53/// The main usage for this is to pass references of the counters across multiple threads while downloading.
54#[derive(Debug, Clone)]
55pub struct ProgressCounter {
56    pub total_mtx: Arc<AtomicUsize>,
57    pub downloaded_mtx: Arc<AtomicU64>,
58    pub main: ProgressBar,
59    pub multi: MultiProgress,
60}
61
62impl ProgressCounter {
63    /// Initialize the main progress bar and the stat counters.
64    ///
65    /// The style that the main progress bar will use is based on the predefined styles for each variant of the ['ImageBoards' enum](ibdl_common::ImageBoards)
66    pub fn initialize(len: u64, imageboard: ImageBoards) -> Self {
67        let template = BarTemplates::new(imageboard);
68        let bar = ProgressBar::new(len).with_style(master_progress_style(&template));
69        bar.set_draw_target(ProgressDrawTarget::stderr());
70        bar.enable_steady_tick(Duration::from_millis(100));
71
72        // Initialize the bars
73        let multi = MultiProgress::new();
74        let main = multi.add(bar);
75
76        Self {
77            main,
78            multi,
79            total_mtx: Default::default(),
80            downloaded_mtx: Default::default(),
81        }
82    }
83
84    /// About the same as `initialize`, but accepts predefined styles instead.
85    pub fn initialize_custom_style(len: u64, style: ProgressStyle) -> Self {
86        let bar = ProgressBar::new(len).with_style(style);
87        bar.set_draw_target(ProgressDrawTarget::stderr());
88        bar.enable_steady_tick(Duration::from_millis(100));
89
90        // Initialize the bars
91        let multi = MultiProgress::new();
92        let main = multi.add(bar);
93
94        Self {
95            main,
96            multi,
97            total_mtx: Default::default(),
98            downloaded_mtx: Default::default(),
99        }
100    }
101
102    /// Adds a download bar under the main progress bar. Will use the predefined style present in the ['ImageBoards' enum](ibdl_common::ImageBoards)
103    pub fn add_download_bar(&self, len: u64, imageboard: ImageBoards) -> ProgressBar {
104        let template = BarTemplates::new(imageboard);
105        let bar = ProgressBar::new(len).with_style(download_progress_style(&template));
106        bar.set_draw_target(ProgressDrawTarget::stderr());
107
108        self.multi.add(bar)
109    }
110
111    /// Adds a download bar under the main progress bar, with a custom style.
112    pub fn add_download_custom_style(&self, len: u64, style: ProgressStyle) -> ProgressBar {
113        let bar = ProgressBar::new(len).with_style(style);
114        bar.set_draw_target(ProgressDrawTarget::stderr());
115
116        self.multi.add(bar)
117    }
118
119    pub async fn init_length_updater(&self, channel: Receiver<u64>) {
120        let mut channel = channel;
121        let cloned_bar = self.main.clone();
122        spawn(async move {
123            while let Some(delta_posts) = channel.recv().await {
124                cloned_bar.inc_length(delta_posts);
125            }
126        })
127        .await
128        .unwrap();
129    }
130
131    pub async fn init_download_counter(&self, channel: Receiver<bool>) {
132        let mut channel = channel;
133        let cloned_bar = self.main.clone();
134        let cloned_mtx = self.downloaded_mtx.clone();
135        spawn(async move {
136            while let Some(downloaded) = channel.recv().await {
137                if downloaded {
138                    cloned_bar.inc(1);
139                    cloned_mtx.fetch_add(1, Ordering::SeqCst);
140                }
141            }
142        });
143    }
144
145    pub fn increment_counters(&self, delta: u64) {
146        self.main.inc(delta);
147        self.total_mtx.fetch_add(delta as usize, Ordering::SeqCst);
148    }
149}
150
151fn master_progress_style(templates: &BarTemplates) -> ProgressStyle {
152    ProgressStyle::default_bar()
153        .template(templates.main)
154        .unwrap()
155        .with_key("pos", |state: &ProgressState, w: &mut dyn Write| {
156            write!(w, "{}/{}", state.pos(), state.len().unwrap()).unwrap();
157        })
158        .with_key("percent", |state: &ProgressState, w: &mut dyn Write| {
159            write!(w, "{:>3.0}%", state.fraction() * 100_f32).unwrap();
160        })
161        .with_key(
162            "files_sec",
163            |state: &ProgressState, w: &mut dyn Write| match state.per_sec() {
164                files_sec if files_sec.abs() < f64::EPSILON => write!(w, "0 files/s").unwrap(),
165                files_sec if files_sec < 1.0 => write!(w, "{:.2} s/file", 1.0 / files_sec).unwrap(),
166                files_sec => write!(w, "{:.2} files/s", files_sec).unwrap(),
167            },
168        )
169        .progress_chars(PROGRESS_CHARS)
170}
171
172fn download_progress_style(templates: &BarTemplates) -> ProgressStyle {
173    ProgressStyle::default_bar()
174        .template(templates.download)
175        .unwrap()
176        .with_key("percent", |state: &ProgressState, w: &mut dyn Write| {
177            write!(w, "{:>3.0}%", state.fraction() * 100_f32).unwrap();
178        })
179        .with_key(
180            "byte_progress",
181            |state: &ProgressState, w: &mut dyn Write| {
182                write!(
183                    w,
184                    "{}/{}",
185                    HumanBytes(state.pos()),
186                    HumanBytes(state.len().unwrap())
187                )
188                .unwrap();
189            },
190        )
191        .progress_chars(PROGRESS_CHARS)
192}
ibdl_core/progress_bars/mod.rs

ibdl_core/progress_bars/
mod.rs
