Allow for local target variables and cli variables

src/main.rs

@@ -3,10 +3,11 @@ extern crate toml;
 
 use clap::Arg;
 use regex::Regex;
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 use std::fs;
 use std::path::Path;
 use std::process::{exit, Command};
+use toml::map::Map;
 
 fn main() {
     //Check if the user has provided a tokefile filename
@@ -26,6 +27,13 @@ fn main() {
                 .required(true)
                 .help("Sets the target to execute"),
         )
+        .arg(
+            Arg::new("variables")
+                .value_delimiter(' ')
+                .num_args(1..)
+                .index(2)
+                .help("User-defined variables in the format KEY=VALUE (these override variables in the tokefile)"),
+        )
         .get_matches();
 
     // Get the target command to execute
@@ -41,7 +49,7 @@ fn main() {
         Some(file_path) => file_path.to_owned(),
         None => {
             let file_names: Vec<String> =
-                vec!["tokefile", "Tokefile", "tokefile.toml", "Tokefile.toml"]
+                ["tokefile", "Tokefile", "tokefile.toml", "Tokefile.toml"]
                     .iter()
                     .map(|s| s.to_string())
                     .collect();
@@ -55,6 +63,7 @@ fn main() {
             tokefile_option.expect("No tokefile found")
         }
     };
+
     let tokefile_contents = match fs::read_to_string(file_path) {
         Ok(contents) => contents.to_string(),
         Err(err) => {
@@ -63,12 +72,25 @@ fn main() {
         }
     };
 
-    // Read the contents of the tokefile
-    //Create all the possible file names
-    //Iterate over the file names and try to read the file
+    let cli_vars = match matches.get_many::<String>("variables") {
+        Some(vars) => vars.collect::<Vec<_>>(),
+        None => vec![],
+    };
+    // Split the variables into a HashMap
+    let cli_vars: HashMap<String, String> = cli_vars
+        .iter()
+        .map(|var| {
+            let parts: Vec<&str> = var.split('=').collect();
+            if parts.len() != 2 {
+                eprintln!("Invalid variable format: {}", var);
+                exit(1);
+            }
+            (parts[0].to_string(), parts[1].to_string())
+        })
+        .collect();
 
     // Parse the tokefile into a TOML Value
-    let parsed_toml = match toml::from_str::<toml::Value>(&tokefile_contents) {
+    let mut parsed_toml = match toml::from_str::<toml::Value>(&tokefile_contents) {
         Ok(value) => value,
         Err(err) => {
             eprintln!("Error parsing tokefile: {}", err);
@@ -76,14 +98,8 @@ fn main() {
         }
     };
 
-    // Extract variables from the tokefile
-    let empty_table = toml::value::Table::new();
-    let binding = parsed_toml.to_owned();
-    let get = binding.get("vars");
-    let vars = get.and_then(|vars| vars.as_table()).unwrap_or(&empty_table);
-
     // Replace variable instances in commands
-    let replaced_commands = replace_variables(parsed_toml.clone(), vars);
+    replace_variables(&mut parsed_toml, cli_vars);
 
     // Determine if there are dependency cycles in the tokefile
     detect_cycle(&parsed_toml);
@@ -91,7 +107,7 @@ fn main() {
     // Check if the target exists
     if parsed_toml
         .get("targets")
-        .and_then(|targets| targets.get(&target).and_then(|t| t.get("cmd")))
+        .and_then(|targets| targets.get(target))
         .is_none()
     {
         eprintln!("Target '{}' not found in tokefile", target);
@@ -99,7 +115,7 @@ fn main() {
     }
 
     // Execute the target command
-    run_command_caller(parsed_toml.clone(), &replaced_commands, target.to_string());
+    run_command(&parsed_toml, target.to_string());
 }
 fn detect_cycle(parsed_toml: &toml::Value) {
     let mut visited_targets = HashSet::new();
@@ -110,7 +126,7 @@ fn detect_cycle(parsed_toml: &toml::Value) {
         .unwrap_or(&empty_table);
 
     for (target_name, _) in targets.iter() {
-        detect_cycle_recursive(&parsed_toml, target_name, &mut visited_targets);
+        detect_cycle_recursive(parsed_toml, target_name, &mut visited_targets);
     }
 }
 
@@ -146,29 +162,63 @@ fn detect_cycle_recursive(
     visited_targets.remove(target);
 }
 
-fn replace_variables(parsed_toml: toml::Value, vars: &toml::value::Table) -> toml::Value {
-    let replaced_commands = parsed_toml
-        .get("targets")
-        .map(|targets| {
-            let mut replaced_targets = toml::value::Table::new();
-            if let Some(targets) = targets.as_table() {
-                for (target_name, target) in targets.iter() {
-                    if let Some(target_table) = target.as_table() {
-                        if let Some(cmd_value) = target_table.get("cmd") {
-                            if let Some(cmd_str) = cmd_value.as_str() {
-                                let replaced_cmd = replace_variables_in_cmd(cmd_str, vars);
-                                replaced_targets
-                                    .insert(target_name.clone(), toml::Value::String(replaced_cmd));
-                            }
-                        }
-                    }
+fn replace_variables(parsed_toml: &mut toml::Value, cli_vars: HashMap<String, String>) {
+    // Parse global variables
+    let map = toml::value::Table::new();
+    let value = &parsed_toml.clone();
+    let get = value.get("vars");
+    let global_vars = get.and_then(|vars| vars.as_table()).unwrap_or(&map);
+
+    // Get the targets table or return an error if it doesn't exist
+    let targets = match parsed_toml
+        .get_mut("targets")
+        .and_then(|targets| targets.as_table_mut())
+    {
+        Some(targets) => targets,
+        None => {
+            eprintln!("No targets found in tokefile");
+            exit(1);
+        }
+    };
+
+    // Iterate over each target
+    for (_, target_value) in targets.iter_mut() {
+        if let Some(target_table) = target_value.as_table_mut() {
+            // Parse local variables for the target
+            let map = toml::value::Table::new();
+            let local_vars = target_table
+                .get("vars")
+                .and_then(|vars| vars.as_table())
+                .unwrap_or(&map);
+
+            // Merge global and local variables
+            let mut merged_vars = merge_vars(global_vars, local_vars);
+
+            // Override variables if they were provided via the CLI
+            for (key, value) in cli_vars.iter() {
+                merged_vars.insert(key.clone(), toml::Value::String(value.clone()));
+            }
+
+            // Replace variables in the target's cmd value
+            if let Some(cmd_value) = target_table.get_mut("cmd") {
+                if let Some(cmd_str) = cmd_value.as_str() {
+                    *cmd_value =
+                        toml::Value::String(replace_variables_in_cmd(cmd_str, &merged_vars));
                 }
             }
-            replaced_targets
-        })
-        .unwrap_or(toml::value::Table::new());
+        }
+    }
+}
 
-    toml::Value::Table(replaced_commands)
+fn merge_vars(
+    global_vars: &Map<String, toml::Value>,
+    local_vars: &Map<String, toml::Value>,
+) -> toml::value::Table {
+    let mut merged_vars = global_vars.clone();
+    for (key, value) in local_vars.iter() {
+        merged_vars.insert(key.clone(), value.clone());
+    }
+    merged_vars
 }
 
 fn replace_variables_in_cmd(cmd: &str, vars: &toml::value::Table) -> String {
@@ -177,7 +227,7 @@ fn replace_variables_in_cmd(cmd: &str, vars: &toml::value::Table) -> String {
     // Regular expression to match variable instances like "${var_name}"
     let re = Regex::new(r#"\$\{([^}]+)\}"#).unwrap();
 
-    for capture in re.captures_iter(&cmd) {
+    for capture in re.captures_iter(cmd) {
         if let Some(var_name) = capture.get(1).map(|m| m.as_str()) {
             if let Some(var_value) = vars.get(var_name).and_then(|v| v.as_str()) {
                 replaced_cmd = replaced_cmd.replace(&format!("${{{}}}", var_name), var_value);
@@ -188,63 +238,39 @@ fn replace_variables_in_cmd(cmd: &str, vars: &toml::value::Table) -> String {
     replaced_cmd
 }
 
-fn run_command_caller(parsed_toml: toml::Value, commands: &toml::Value, target: String) {
-    //Create a hashset to keep track of visited targets
-    let visited_targets = HashSet::new();
-    run_command(parsed_toml, commands, target, &visited_targets);
-}
-
-fn run_command(
-    parsed_toml: toml::Value,
-    commands: &toml::Value,
-    target: String,
-    visited_targets: &HashSet<String>,
-) {
-    //Check if the target exists
-    match commands.get(target.clone()) {
-        Some(some_target) => {
-            //Execute it's dependencies first, by order of appearance
-            if let Some(target) = parsed_toml
-                .get("targets")
-                .and_then(|targets| targets.get(&target))
-            {
-                if let Some(target_table) = target.as_table() {
-                    if let Some(dep_value) = target_table.get("deps") {
-                        if let Some(dep_array) = dep_value.as_array() {
-                            for dep in dep_array {
-                                if let Some(dep_str) = dep.as_str() {
-                                    run_command(
-                                        parsed_toml.clone(),
-                                        commands,
-                                        dep_str.to_string(),
-                                        visited_targets,
-                                    );
-                                }
-                            }
+fn run_command(parsed_toml: &toml::Value, target: String) {
+    if let Some(targets) = parsed_toml
+        .get("targets")
+        .and_then(|targets| targets.as_table())
+    {
+        if let Some(target_table) = targets.get(&target) {
+            if let Some(dep_value) = target_table.get("deps") {
+                if let Some(dep_array) = dep_value.as_array() {
+                    for dep in dep_array {
+                        if let Some(dep_str) = dep.as_str() {
+                            run_command(parsed_toml, dep_str.to_string());
                         }
                     }
                 }
             }
-            if let Some(cmd) = some_target.as_str() {
-                eprintln!("{}", cmd);
-                let status = Command::new("sh")
-                    .arg("-c")
-                    .arg(cmd)
-                    .status()
-                    .expect("Failed to execute command");
-                if !status.success() {
-                    eprintln!(
-                        "Command '{}' failed with exit code {:?}",
-                        cmd,
-                        status.code()
-                    );
-                    exit(1);
+            if let Some(cmd_value) = target_table.get("cmd") {
+                if let Some(cmd_str) = cmd_value.as_str() {
+                    eprintln!("{}", cmd_str);
+                    let status = Command::new("sh")
+                        .arg("-c")
+                        .arg(cmd_str)
+                        .status()
+                        .expect("Failed to execute command");
+                    if !status.success() {
+                        eprintln!(
+                            "Command '{}' failed with exit code {:?}",
+                            cmd_str,
+                            status.code()
+                        );
+                        exit(1);
+                    }
                 }
             }
         }
-        None => {
-            eprintln!("Target '{}' not found in tokefile", target);
-            exit(1);
-        }
     }
 }