Skip to main content

zenodo_rs/
progress.rs

1//! Progress reporting hooks for uploads and downloads.
2//!
3//! The core client stays generic over any progress sink that implements
4//! [`TransferProgress`]. Under the optional `indicatif` feature,
5//! [`indicatif::ProgressBar`] implements this trait directly.
6//!
7//! # Examples
8//!
9//! ```rust
10//! #[cfg(feature = "indicatif")]
11//! #[tokio::main]
12//! async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
13//!     use axum::{
14//!         body::Body,
15//!         extract::State,
16//!         http::{header, HeaderValue},
17//!         routing::get,
18//!         Json, Router,
19//!     };
20//!     use indicatif::{ProgressBar, ProgressStyle};
21//!     use serde_json::json;
22//!     use std::sync::Arc;
23//!     use zenodo_rs::{ArtifactSelector, Auth, Endpoint, RecordId, ZenodoClient};
24//!
25//!     #[derive(Clone)]
26//!     struct AppState {
27//!         base: Arc<String>,
28//!     }
29//!
30//!     async fn record(State(state): State<AppState>) -> Json<serde_json::Value> {
31//!         Json(json!({
32//!             "id": 123,
33//!             "recid": 123,
34//!             "metadata": { "title": "Example" },
35//!             "files": [{
36//!                 "id": "f-123",
37//!                 "key": "artifact.bin",
38//!                 "size": 5,
39//!                 "links": {
40//!                     "self": format!("{}download/123/artifact.bin", state.base),
41//!                 }
42//!             }],
43//!             "links": {}
44//!         }))
45//!     }
46//!
47//!     async fn artifact() -> axum::response::Response {
48//!         let mut response = axum::response::Response::new(Body::from("hello"));
49//!         response.headers_mut().insert(
50//!             header::CONTENT_TYPE,
51//!             HeaderValue::from_static("application/octet-stream"),
52//!         );
53//!         response
54//!     }
55//!
56//!     let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
57//!     let address = listener.local_addr()?;
58//!     let base = Arc::new(format!("http://{address}/api/"));
59//!     let app = Router::new()
60//!         .route("/api/records/123", get(record))
61//!         .route("/api/download/123/artifact.bin", get(artifact))
62//!         .with_state(AppState { base: Arc::clone(&base) });
63//!     let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>();
64//!     let server = tokio::spawn(async move {
65//!         axum::serve(listener, app)
66//!             .with_graceful_shutdown(async {
67//!                 let _ = shutdown_rx.await;
68//!             })
69//!             .await
70//!     });
71//!
72//!     let client = ZenodoClient::builder(Auth::new("token"))
73//!         .endpoint(Endpoint::Custom(base.parse()?))
74//!         .build()?;
75//!     let bar = ProgressBar::new(0);
76//!     bar.set_style(ProgressStyle::with_template(
77//!         "{bar:20.cyan/blue} {bytes}/{total_bytes}",
78//!     )?);
79//!     let temp_dir = tempfile::tempdir()?;
80//!     let path = temp_dir.path().join("artifact.bin");
81//!
82//!     let resolved = client
83//!         .download_artifact_with_progress(
84//!             &ArtifactSelector::latest_file(RecordId(123), "artifact.bin"),
85//!             &path,
86//!             bar.clone(),
87//!         )
88//!         .await?;
89//!
90//!     assert_eq!(resolved.bytes_written, 5);
91//!     assert_eq!(std::fs::read(&path)?, b"hello");
92//!     assert_eq!(bar.position(), 5);
93//!     let _ = shutdown_tx.send(());
94//!     server.await??;
95//!     Ok(())
96//! }
97//!
98//! #[cfg(not(feature = "indicatif"))]
99//! fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
100//!     Ok(())
101//! }
102//! ```
103//!
104//! Pass `bar.clone()` into the progress-aware upload and download helpers when
105//! you want a real terminal progress bar during transfers.
106
107/// Progress sink for streaming uploads and downloads.
108///
109/// Implement this trait when you want upload and download helpers to report
110/// byte-level transfer progress into your own logging, UI, or terminal
111/// progress bar implementation.
112pub trait TransferProgress: Send + Sync {
113    /// Called once before the transfer starts.
114    ///
115    /// `total_bytes` is `Some(len)` when the total size is known up front and
116    /// `None` when the transfer length is unknown.
117    fn begin(&self, _total_bytes: Option<u64>) {}
118
119    /// Called after each successfully transferred chunk.
120    fn advance(&self, _delta: u64) {}
121
122    /// Called once after a transfer completes successfully.
123    fn finish(&self) {}
124}
125
126impl TransferProgress for () {}
127
128impl<P> TransferProgress for std::sync::Arc<P>
129where
130    P: TransferProgress + ?Sized,
131{
132    fn begin(&self, total_bytes: Option<u64>) {
133        self.as_ref().begin(total_bytes);
134    }
135
136    fn advance(&self, delta: u64) {
137        self.as_ref().advance(delta);
138    }
139
140    fn finish(&self) {
141        self.as_ref().finish();
142    }
143}
144
145#[cfg(feature = "indicatif")]
146impl TransferProgress for indicatif::ProgressBar {
147    fn begin(&self, total_bytes: Option<u64>) {
148        self.set_position(0);
149        if let Some(total_bytes) = total_bytes {
150            self.set_length(total_bytes);
151        }
152    }
153
154    fn advance(&self, delta: u64) {
155        self.inc(delta);
156    }
157
158    fn finish(&self) {
159        indicatif::ProgressBar::finish(self);
160    }
161}
162
163#[cfg(all(test, feature = "indicatif"))]
164mod tests {
165    use super::TransferProgress;
166
167    #[test]
168    fn indicatif_progress_bar_tracks_transfer_progress() {
169        let bar = indicatif::ProgressBar::new(0);
170
171        bar.begin(Some(5));
172        bar.advance(2);
173
174        assert_eq!(bar.length(), Some(5));
175        assert_eq!(bar.position(), 2);
176
177        TransferProgress::finish(&bar);
178        assert!(bar.is_finished());
179    }
180}