hf_fetch_model/progress.rs
1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! Progress reporting for model downloads.
4//!
5//! [`ProgressEvent`] carries per-file and overall download status.
6//! When the `indicatif` feature is enabled, `IndicatifProgress`
7//! provides multi-progress bars out of the box.
8
9/// A progress event emitted during download.
10///
11/// Passed to the `on_progress` callback on [`crate::FetchConfig`].
12#[derive(Debug, Clone)]
13pub struct ProgressEvent {
14 /// The filename currently being downloaded.
15 pub filename: String,
16 /// Bytes downloaded so far for this file.
17 pub bytes_downloaded: u64,
18 /// Total size of this file in bytes (0 if unknown).
19 pub bytes_total: u64,
20 /// Download percentage for this file (0.0–100.0).
21 pub percent: f64,
22 /// Number of files still remaining (after this one).
23 pub files_remaining: usize,
24}
25
26/// Creates a [`ProgressEvent`] for a completed file.
27#[must_use]
28pub(crate) fn completed_event(filename: &str, size: u64, files_remaining: usize) -> ProgressEvent {
29 ProgressEvent {
30 filename: filename.to_owned(),
31 bytes_downloaded: size,
32 bytes_total: size,
33 percent: 100.0,
34 files_remaining,
35 }
36}
37
38/// Multi-progress bar display using `indicatif`.
39///
40/// Available only when the `indicatif` feature is enabled.
41///
42/// # Example
43///
44/// ```rust,no_run
45/// # fn example() -> Result<(), hf_fetch_model::FetchError> {
46/// use hf_fetch_model::FetchConfig;
47/// # #[cfg(feature = "indicatif")]
48/// use hf_fetch_model::progress::IndicatifProgress;
49///
50/// # #[cfg(feature = "indicatif")]
51/// let progress = IndicatifProgress::new();
52/// let config = FetchConfig::builder()
53/// # ;
54/// # #[cfg(feature = "indicatif")]
55/// # let config = FetchConfig::builder()
56/// .on_progress(move |e| progress.handle(e))
57/// .build()?;
58/// # Ok(())
59/// # }
60/// ```
61#[cfg(feature = "indicatif")]
62pub struct IndicatifProgress {
63 multi: indicatif::MultiProgress,
64 overall: indicatif::ProgressBar,
65}
66
67#[cfg(feature = "indicatif")]
68impl IndicatifProgress {
69 /// Creates a new multi-progress bar display.
70 ///
71 /// Call [`IndicatifProgress::set_total_files`] once the file count is known.
72 #[must_use]
73 pub fn new() -> Self {
74 let multi = indicatif::MultiProgress::new();
75 let overall = multi.add(indicatif::ProgressBar::new(0));
76 overall.set_style(
77 indicatif::ProgressStyle::default_bar()
78 .template("{msg} [{bar:40.cyan/blue}] {pos}/{len} files")
79 .ok()
80 .unwrap_or_else(indicatif::ProgressStyle::default_bar)
81 .progress_chars("=> "),
82 );
83 overall.set_message("Overall");
84 Self { multi, overall }
85 }
86
87 /// Returns a reference to the underlying [`indicatif::MultiProgress`].
88 ///
89 /// Useful for adding custom progress bars alongside the built-in ones.
90 #[must_use]
91 pub fn multi(&self) -> &indicatif::MultiProgress {
92 &self.multi
93 }
94
95 /// Sets the total number of files to download.
96 pub fn set_total_files(&self, total: u64) {
97 self.overall.set_length(total);
98 }
99
100 /// Handles a [`ProgressEvent`], updating progress bars.
101 pub fn handle(&self, event: &ProgressEvent) {
102 if event.percent >= 100.0 {
103 // Derive total: completed so far + this file + remaining
104 // EXPLICIT: try_from for usize → u64 (infallible on 64-bit, safe fallback otherwise)
105 let remaining = u64::try_from(event.files_remaining).unwrap_or(u64::MAX);
106 let total = self.overall.position() + 1 + remaining;
107 self.overall.set_length(total);
108 self.overall.inc(1);
109 }
110 }
111}
112
113#[cfg(feature = "indicatif")]
114impl Default for IndicatifProgress {
115 fn default() -> Self {
116 Self::new()
117 }
118}