tiny_data/
client.rs

1use futures::future::join_all;
2use indicatif::ProgressBar;
3use indicatif::{MultiProgress, ProgressStyle};
4use std::fs::create_dir_all;
5use std::sync::Arc;
6use std::time::Instant;
7use tokio::task::JoinHandle;
8
9use crate::download::*;
10use crate::fetch::*;
11
12struct TopicHandler<'a> {
13    topic: &'a str,
14    download_manager: DLManager,
15}
16
17impl<'a> TopicHandler<'a> {
18    pub async fn download_topic(&mut self, dir: String) -> u8 {
19        let nsamples = self.download_manager.target_size;
20
21        //download path = base_dir + topic
22        let dir = format!("{}/{}", dir, self.topic);
23
24        let batch = Fetcher::query_api(self.topic, nsamples).await;
25        let total = self.download_manager.download_batch(batch, &dir).await;
26        total
27    }
28}
29
30pub struct TinyDataClient {
31    multi_progress_bar: MultiProgress,
32}
33
34impl TinyDataClient {
35    pub fn new() -> Self {
36        Self {
37            multi_progress_bar: MultiProgress::new(),
38        }
39    }
40
41    pub async fn run(&self, topics: Vec<String>, nsamples: usize, dir: String) {
42        //create image directories
43        for topic in &topics {
44            let dir = format!("{}/{}", dir, topic);
45            create_dir_all(&dir).expect("failed to create directory");
46        }
47
48        // need thread-safe string for async
49        let dir = Arc::new(dir);
50        let topics_len = topics.len();
51
52        let futures: Vec<JoinHandle<u8>> = topics
53            .into_iter()
54            .map(|topic| {
55                let dir = Arc::clone(&dir);
56                let mut pb = self
57                    .multi_progress_bar
58                    .add(ProgressBar::new(nsamples as u64));
59
60                //style
61                stylize_pb(&mut pb, &topic);
62
63                tokio::spawn(async move {
64                    let download_manager = DLManager::new(nsamples, pb);
65
66                    let mut topic_handler = TopicHandler {
67                        topic: &topic,
68                        download_manager,
69                    };
70                    topic_handler.download_topic(dir.to_string()).await
71                })
72            })
73            .collect();
74
75        //time execution
76        let now = Instant::now();
77        let total = join_all(futures).await;
78        let elapsed = now.elapsed();
79
80        let total: u8 = total.into_iter().map(|res| res.unwrap()).sum();
81
82        println!(
83            "{}/{} files saved successfully to `./{}` in {}s 📦",
84            total,
85            nsamples * topics_len,
86            dir,
87            elapsed.as_secs(),
88        );
89    }
90}
91
92fn stylize_pb(pb: &mut ProgressBar, name: &str) {
93    let default = "{msg} {spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})";
94    pb.set_style(
95        ProgressStyle::with_template(default)
96            .unwrap()
97            .progress_chars("#>-"),
98    );
99    pb.set_message(String::from(name));
100}
tiny_data/client.rs

tiny_data/
client.rs
