Add input variables and improve parsing

Signed-off-by: trivernis <trivernis@protonmail.com>

Cargo.lock

@@ -3281,6 +3281,8 @@ dependencies = [
  "miette 5.3.0",
  "nu-cli",
  "nu-command",
+ "nu-engine",
+ "nu-parser",
  "nu-protocol",
  "thiserror",
  "tokio",

Cargo.toml

@@ -19,6 +19,8 @@ lazy_static = "1.4.0"
 miette = "5.3.0"
 nu-cli = "0.68.1"
 nu-command = "0.68.1"
+nu-engine = "0.68.1"
+nu-parser = "0.68.1"
 nu-protocol = "0.68.1"
 thiserror = "1.0.35"
 tokio = { version = "1.21.1", features = ["rt", "io-std", "io-util", "process", "time", "macros", "tracing", "fs"] }

src/error.rs

@@ -13,6 +13,12 @@ pub enum AppError {
 
     #[error("Could not find the script file {0}")]
     ScriptNotFound(PathBuf),
+
+    #[error("Could not parse the source file {0}")]
+    ParseError(#[from] nu_parser::ParseError),
+
+    #[error("Could not find the main mehod in the script file {0}")]
+    MissingMain(PathBuf),
 }
 
 impl From<miette::Error> for AppError {

src/scripting/executor.rs

@@ -1,14 +1,38 @@
 #![allow(unused)]
-use std::path::Path;
+use std::{
+    collections::HashMap,
+    fs, mem,
+    path::{Path, PathBuf},
+};
 
-use nu_protocol::{PipelineData, Span};
+use miette::IntoDiagnostic;
+use nu_protocol::{
+    ast::{Block, Call},
+    engine::{EngineState, Stack, StateDelta, StateWorkingSet},
+    PipelineData, Span, Type, Value,
+};
 
 use crate::error::{AppError, AppResult};
 
+#[derive(Clone, Debug)]
+pub enum VarValue {
+    String(String),
+    Integer(i64),
+    Float(f64),
+    Boolean(bool),
+}
+
+impl VarValue {
+    pub fn string<S: ToString>(s: S) -> Self {
+        Self::String(s.to_string())
+    }
+}
+
 /// An executor for nu scripts
 pub struct NuExecutor {
     script_path: String,
     args: Vec<String>,
+    global_vars: HashMap<String, VarValue>,
 }
 
 impl NuExecutor {
@@ -16,6 +40,7 @@ impl NuExecutor {
         Self {
             script_path: script_path.as_ref().to_string_lossy().into_owned(),
             args: Vec::new(),
+            global_vars: HashMap::new(),
         }
     }
 
@@ -32,21 +57,102 @@ impl NuExecutor {
         self
     }
 
+    /// Adds a global variable to the executor which can
+    /// be accessed from within the script
+    pub fn add_global_var<S: ToString>(&mut self, name: S, value: VarValue) -> &mut Self {
+        self.global_vars.insert(name.to_string(), value);
+
+        self
+    }
+
+    /// Executes the given script file in a clean nu context.
     pub fn execute(&mut self) -> AppResult<()> {
         let mut engine_state = nu_command::create_default_context();
         let mut stack = nu_protocol::engine::Stack::new();
         let input = PipelineData::new(Span::new(0, 0));
         let init_cwd = nu_cli::get_init_cwd();
-        nu_cli::gather_parent_env_vars(&mut engine_state, &init_cwd);
+        nu_engine::convert_env_values(&mut engine_state, &mut stack);
 
-        nu_cli::evaluate_file(
-            self.script_path.clone(),
-            &self.args,
-            &mut engine_state,
+        let vars = mem::take(&mut self.global_vars);
+        let vars = vars
+            .into_iter()
+            .map(|(k, v)| (k, map_var_to_value(v)))
+            .collect::<HashMap<_, _>>();
+
+        add_variables_to_state(vars, &mut engine_state, &mut stack);
+        let block = read_script_file(&self.script_path, &mut engine_state)?;
+
+        nu_engine::eval_block(
+            &engine_state,
             &mut stack,
-            input,
+            &block,
+            PipelineData::new(Span::new(0, 0)),
+            false,
             false,
         )
-        .map_err(AppError::from)
+        .into_diagnostic()?;
+
+        // TODO: Create the AST for the call here instead of parsing it from a string
+        let args = format!("main {}", self.args.join(" "));
+        nu_cli::eval_source(
+            &mut engine_state,
+            &mut stack,
+            args.as_bytes(),
+            "<commandline>",
+            input,
+        );
+
+        Ok(())
+    }
+}
+
+/// Adds variables to the nu engine state
+/// Note: Calling this function multiple times will override other variables
+fn add_variables_to_state(
+    vars: HashMap<String, Value>,
+    state: &mut EngineState,
+    stack: &mut Stack,
+) {
+    let state2 = nu_command::create_default_context();
+    let mut working_set = StateWorkingSet::new(&state2);
+    for (name, value) in vars {
+        let var_id = working_set.add_variable(
+            name.as_bytes().to_vec(),
+            Span::new(0, 0),
+            nu_protocol::Type::String,
+        );
+        stack.add_var(var_id, value);
+    }
+    state.merge_delta(working_set.render());
+}
+
+/// Reads the nu script file and
+/// returns its root block
+fn read_script_file(path: &str, engine_state: &mut EngineState) -> AppResult<Block> {
+    let mut working_set = StateWorkingSet::new(engine_state);
+    let script_contents = fs::read(&path).into_diagnostic()?;
+    // parse the source file
+    let (block, err) = nu_parser::parse(&mut working_set, Some(path), &script_contents, false, &[]);
+
+    if let Some(err) = err {
+        return Err(AppError::from(err));
+    }
+    // check if a main method exists in the block
+    if !working_set.find_decl(b"main", &Type::Block).is_some() {
+        return Err(AppError::MissingMain(PathBuf::from(path)));
+    }
+    engine_state.merge_delta(working_set.render());
+
+    Ok(block)
+}
+
+fn map_var_to_value(var: VarValue) -> Value {
+    let span = Span::new(0, 0);
+
+    match var {
+        VarValue::String(val) => Value::String { val, span },
+        VarValue::Integer(val) => Value::Int { val, span },
+        VarValue::Float(val) => Value::Float { val, span },
+        VarValue::Boolean(val) => Value::Bool { val, span },
     }
 }

src/scripting/script.rs

@@ -2,7 +2,7 @@ use std::{marker::PhantomData, path::PathBuf};
 
 use crate::error::AppResult;
 
-use super::executor::NuExecutor;
+use super::executor::{NuExecutor, VarValue};
 
 /// A trait implemented for a given nu script type to
 /// associate arguments
@@ -39,6 +39,11 @@ impl<S: Script> NuScript<S> {
     pub fn execute(&self, args: S::Args) -> AppResult<()> {
         NuExecutor::new(&self.path)
             .add_args(args.get_args())
+            .add_global_var("BY_TOURMALINE", VarValue::string("Hello from Tourmaline!"))
+            .add_global_var(
+                "ANOTHER_ONE",
+                VarValue::string("This variable was provided by tourmaline"),
+            )
             .execute()
     }
 }