From ad8e831a5cbe5bb65d0567d4bdde95bd3ef3de75 Mon Sep 17 00:00:00 2001
From: Lia Lenckowski <lialenck@protonmail.com>
Date: Thu, 7 Sep 2023 16:41:02 +0200
Subject: add progress bars/spinners

---
 src/ai_embedders.rs | 25 +++++++++++++++++++++----
 1 file changed, 21 insertions(+), 4 deletions(-)

(limited to 'src/ai_embedders.rs')

diff --git a/src/ai_embedders.rs b/src/ai_embedders.rs
index 14e1d8c..b30d890 100644
--- a/src/ai_embedders.rs
+++ b/src/ai_embedders.rs
@@ -1,6 +1,7 @@
 use anyhow::Result;
+use indicatif::{ProgressBar, ProgressIterator, ProgressStyle};
 use serde::{Deserialize, Serialize};
-use std::{io::{copy, Cursor}, path::PathBuf, process::Command};
+use std::{io::{BufRead, BufReader, copy, Cursor}, path::PathBuf, process::{Command, Stdio}};
 
 use crate::{Config, BatchEmbedder, MetricElem};
 
@@ -46,20 +47,36 @@ impl BatchEmbedder for ContentEmbedder<'_> {
         let api_prog = include_bytes!("imgbeddings-api.py");
         copy(&mut Cursor::new(api_prog), &mut std::fs::File::create(&script_file)?)?;
 
+        let bar = ProgressBar::new_spinner();
+        bar.set_style(ProgressStyle::with_template("{spinner} {msg}")?);
+        bar.enable_steady_tick(std::time::Duration::from_millis(100));
+
+        bar.set_message("Creating venv...");
         Command::new("python3")
             .args(["-m", "venv", venv_dir.to_str().unwrap()])
+            .stdout(Stdio::null())
             .spawn()?
             .wait()?;
+
+        bar.set_message("Installing/checking packages...");
         Command::new(venv_dir.join("bin/pip3"))
             .args(["install", "imgbeddings"])
+            .stdout(Stdio::null())
             .spawn()?
             .wait()?;
+        bar.finish();
 
-        let output = Command::new(venv_dir.join("bin/python3"))
+        let child = Command::new(venv_dir.join("bin/python3"))
             .arg(script_file)
             .args(paths)
-            .output()?;
+            .stderr(Stdio::null())
+            .stdout(Stdio::piped())
+            .spawn()?;
 
-        Ok(serde_json::from_slice(&output.stdout)?)
+        Ok(BufReader::new(child.stdout.unwrap())
+            .lines()
+            .progress_count(paths.len().try_into().unwrap())
+            .map(|l| Ok::<_, anyhow::Error>(serde_json::from_str(&l?)?))
+            .try_collect()?)
     }
 }
-- 
cgit v1.2.3-70-g09d2