From 5ed0246e2bf0804c8d18f795cfcf5f80d7ad5fab Mon Sep 17 00:00:00 2001
From: Couleur <82747632+couleurm@users.noreply.github.com>
Date: Thu, 22 Feb 2024 10:33:14 +0100
Subject: [PATCH] add --variables, -v to set list of key/value script env
 variables

---
 src/main.rs           | 10 +++++++++-
 src/vs_handler/mod.rs | 21 ++++++++++++++++++++-
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 78cada5..baa82ed 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -18,6 +18,14 @@ use vs_handler::PreviewedScript;
 struct Opt {
     #[arg(id = "input", value_hint = ValueHint::FilePath)]
     input: PathBuf,
+
+    #[arg(
+        id = "variable",
+        short = 'v',
+        help = "Variables to set in the script environment. Example: `-v key=value`",
+        value_delimiter = ','
+    )]
+    variables: Vec<String>,
 }
 
 #[tokio::main(flavor = "multi_thread", worker_threads = 1)]
@@ -28,7 +36,7 @@ async fn main() -> Result<()> {
         bail!("Input script file does not exist!");
     }
 
-    let script = Arc::new(Mutex::new(PreviewedScript::new(opt.input)));
+    let script = Arc::new(Mutex::new(PreviewedScript::new(opt.input, opt.variables)));
     let (cmd_sender, cmd_receiver) = tokio::sync::mpsc::channel(1);
 
     {
diff --git a/src/vs_handler/mod.rs b/src/vs_handler/mod.rs
index dfd2e21..6675ddf 100644
--- a/src/vs_handler/mod.rs
+++ b/src/vs_handler/mod.rs
@@ -1,3 +1,4 @@
+use itertools::Itertools;
 use parking_lot::Mutex;
 use std::collections::HashMap;
 use std::path::PathBuf;
@@ -22,6 +23,7 @@ pub use vstransform::*;
 pub struct PreviewedScript {
     script_file: String,
     script_dir: PathBuf,
+    variables: Vec<String>,
 
     #[serde(skip)]
     env: Option<Environment>,
@@ -45,7 +47,7 @@ pub struct VSMessage {
 }
 
 impl PreviewedScript {
-    pub fn new(script_path: PathBuf) -> Self {
+    pub fn new(script_path: PathBuf, variables: Vec<String>) -> Self {
         let mut script_dir = script_path.clone();
         script_dir.pop();
 
@@ -57,6 +59,7 @@ impl PreviewedScript {
         Self {
             script_file,
             script_dir,
+            variables,
             env: None,
             message_handler_id: None,
             vs_messages: Arc::new(Mutex::new(Vec::new())),
@@ -89,6 +92,22 @@ impl PreviewedScript {
             self.message_handler_id = Some(id);
         }
 
+        if !self.variables.is_empty() {
+            let mut variables = OwnedMap::new(API::get().unwrap());
+            let kv_map = self
+                .variables
+                .iter()
+                .filter_map(|s| s.split('=').collect_tuple());
+            for (name, value) in kv_map {
+                variables
+                    .append_data(name, value.as_bytes())
+                    .expect("Couldn't append an argument value");
+            }
+
+            env.set_variables(&variables)
+                .expect("Couldn't set arguments");
+        }
+
         env.eval_file(&self.script_file, EvalFlags::SetWorkingDir)?;
 
         Ok(())