Add AST builder for script arguments

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

src/lib.rs

@@ -1,5 +1,5 @@
 use error::AppResult;
-use scripting::{loader::ScriptLoader, script::JSONArgs};
+use scripting::loader::ScriptLoader;
 use tasks::{SetupUsersScript, UsersConfig};
 
 pub mod error;
@@ -12,18 +12,20 @@ pub struct TaskExecutor {
 }
 
 impl TaskExecutor {
-    pub fn new() -> Self {
-        Self {
-            loader: ScriptLoader::new(),
-        }
-    }
-
     /// Sets up user accounts
     #[tracing::instrument(level = "debug", skip(self))]
     pub async fn setup_users(&self, users_cfg: UsersConfig) -> AppResult<()> {
         self.loader
             .load::<SetupUsersScript>()?
-            .execute(JSONArgs(users_cfg))
+            .execute(users_cfg)
             .await
     }
 }
+
+impl Default for TaskExecutor {
+    fn default() -> Self {
+        Self {
+            loader: ScriptLoader::new(),
+        }
+    }
+}

src/main.rs

@@ -7,7 +7,7 @@ use tourmaline::{
 async fn main() -> miette::Result<()> {
     color_eyre::install().unwrap();
     dotenv::dotenv().unwrap();
-    let executor = TaskExecutor::new();
+    let executor = TaskExecutor::default();
     let user_cfg = UsersConfig {
         users: vec![
             User {

src/scripting/executor.rs

@@ -10,13 +10,15 @@ use tokio::fs;
 
 use miette::IntoDiagnostic;
 use nu_protocol::{
-    ast::{Block, Call},
+    ast::{Argument, Block, Call, Expr, Expression, Pipeline},
     engine::{EngineState, Stack, StateDelta, StateWorkingSet},
-    PipelineData, Span, Type, Value,
+    DeclId, PipelineData, Signature, Span, Type, Value,
 };
 
 use crate::error::{AppError, AppResult};
 
+use super::record::RecordValue;
+
 #[derive(Clone, Debug)]
 pub enum VarValue {
     String(String),
@@ -34,7 +36,7 @@ impl VarValue {
 /// An executor for nu scripts
 pub struct NuExecutor {
     script_path: PathBuf,
-    args: Vec<String>,
+    args: Vec<RecordValue>,
     global_vars: HashMap<String, VarValue>,
 }
 
@@ -47,14 +49,17 @@ impl NuExecutor {
         }
     }
 
-    pub fn add_arg<S: ToString>(&mut self, arg: S) -> &mut Self {
-        self.args.push(arg.to_string());
+    pub fn add_arg<A: Into<RecordValue>>(&mut self, arg: A) -> &mut Self {
+        self.args.push(arg.into());
 
         self
     }
 
-    pub fn add_args<S: ToString, I: IntoIterator<Item = S>>(&mut self, args: I) -> &mut Self {
-        let mut args = args.into_iter().map(|a| a.to_string()).collect::<Vec<_>>();
+    pub fn add_args<A: Into<RecordValue>, I: IntoIterator<Item = A>>(
+        &mut self,
+        args: I,
+    ) -> &mut Self {
+        let mut args = args.into_iter().map(|a| a.into()).collect::<Vec<_>>();
         self.args.append(&mut args);
 
         self
@@ -74,7 +79,7 @@ impl NuExecutor {
         let mut stack = nu_protocol::engine::Stack::new();
         let init_cwd = nu_cli::get_init_cwd();
         gather_parent_env_vars(&mut engine_state, &init_cwd);
-        nu_engine::convert_env_values(&mut engine_state, &mut stack);
+        nu_engine::convert_env_values(&mut engine_state, &stack);
 
         let vars = mem::take(&mut self.global_vars);
         let vars = vars
@@ -83,7 +88,7 @@ impl NuExecutor {
             .collect::<HashMap<_, _>>();
 
         add_variables_to_state(vars, &mut engine_state, &mut stack);
-        let block = read_script_file(&self.script_path, &mut engine_state).await?;
+        let (block, main_id) = read_script_file(&self.script_path, &mut engine_state).await?;
 
         nu_engine::eval_block(
             &engine_state,
@@ -99,11 +104,7 @@ impl NuExecutor {
 
         // block in a different thread to be able to execute scripts in parallel
         tokio::task::spawn_blocking(move || {
-            // TODO: Create the AST for the call here instead of parsing it from a string
-            let args = format!("main {}", args.join(" "));
-            let (call_block, working_set) = parse_nu(&mut engine_state, args.as_bytes(), None)?;
-            let delta = working_set.render();
-            engine_state.merge_delta(delta);
+            let call_block = create_call(main_id, args);
 
             nu_engine::eval_block(
                 &engine_state,
@@ -149,19 +150,23 @@ fn add_variables_to_state(
 
 /// Reads the nu script file and
 /// returns its root block
-async fn read_script_file(path: &Path, engine_state: &mut EngineState) -> AppResult<Block> {
+async fn read_script_file(
+    path: &Path,
+    engine_state: &mut EngineState,
+) -> AppResult<(Block, DeclId)> {
     let script_contents = fs::read(&path).await.into_diagnostic()?;
     let string_path = path.to_string_lossy().into_owned();
     // parse the source file
     let (block, working_set) = parse_nu(engine_state, &script_contents, Some(&string_path))?;
     // 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)));
-    }
-    let delta = working_set.render();
-    engine_state.merge_delta(delta);
+    if let Some(decl_id) = working_set.find_decl(b"main", &Type::Block) {
+        let delta = working_set.render();
+        engine_state.merge_delta(delta);
 
-    Ok(block)
+        Ok((block, decl_id))
+    } else {
+        Err(AppError::MissingMain(PathBuf::from(path)))
+    }
 }
 
 fn parse_nu<'a>(
@@ -179,6 +184,35 @@ fn parse_nu<'a>(
     }
 }
 
+fn create_call(decl_id: DeclId, args: Vec<RecordValue>) -> Block {
+    let args = args
+        .into_iter()
+        .map(|a| Argument::Positional(a.into_expression()))
+        .collect();
+    let call = Call {
+        decl_id,
+        head: Span::new(0, 0),
+        arguments: args,
+        redirect_stdout: true,
+        redirect_stderr: false,
+    };
+    let pipeline = Pipeline {
+        expressions: vec![Expression {
+            expr: Expr::Call(Box::new(call)),
+            span: Span::new(0, 0),
+            ty: Type::Any,
+            custom_completion: None,
+        }],
+    };
+    Block {
+        signature: Box::new(Signature::build("Call to main")),
+        pipelines: vec![pipeline],
+        captures: Vec::new(),
+        redirect_env: false,
+        span: None,
+    }
+}
+
 fn empty_pipeline() -> PipelineData {
     PipelineData::new(Span::new(0, 0))
 }

src/scripting/mod.rs

@@ -1,3 +1,4 @@
 pub mod executor;
 pub mod loader;
+pub mod record;
 pub mod script;

src/scripting/record.rs

@@ -0,0 +1,147 @@
+use std::{borrow::Cow, collections::HashMap};
+
+use nu_protocol::{
+    ast::{Expr, Expression},
+    Span, Type,
+};
+
+#[derive(Clone, Debug)]
+pub enum RecordValue {
+    Int(i64),
+    Float(f64),
+    String(String),
+    Boolean(bool),
+    Null,
+    Map(Vec<(RecordValue, RecordValue)>),
+    List(Vec<RecordValue>),
+}
+
+impl RecordValue {
+    pub fn into_expr(self) -> Expr {
+        match self {
+            RecordValue::Int(i) => Expr::Int(i),
+            RecordValue::Float(f) => Expr::Float(f),
+            RecordValue::String(s) => Expr::String(s),
+            RecordValue::Boolean(b) => Expr::Bool(b),
+            RecordValue::Null => Expr::Nothing,
+            RecordValue::Map(m) => Expr::Record(
+                m.into_iter()
+                    .map(|(k, v)| (k.into_expression(), v.into_expression()))
+                    .collect(),
+            ),
+            RecordValue::List(l) => {
+                Expr::List(l.into_iter().map(RecordValue::into_expression).collect())
+            }
+        }
+    }
+
+    pub fn into_expression(self) -> Expression {
+        let nu_type = self.get_type();
+        let expr = self.into_expr();
+
+        Expression {
+            expr,
+            span: Span::new(0, 0),
+            ty: nu_type,
+            custom_completion: None,
+        }
+    }
+
+    fn get_type(&self) -> Type {
+        match &self {
+            RecordValue::Int(_) => Type::Int,
+            RecordValue::Float(_) => Type::Float,
+            RecordValue::String(_) => Type::String,
+            RecordValue::Boolean(_) => Type::Bool,
+            RecordValue::Null => Type::Nothing,
+            RecordValue::Map(m) => {
+                let type_map = m
+                    .iter()
+                    .map(|(k, v)| (k.to_string(), v.get_type()))
+                    .collect();
+                Type::Record(type_map)
+            }
+            RecordValue::List(l) => {
+                let list_type = if let Some(first) = l.first() {
+                    first.get_type()
+                } else {
+                    Type::Nothing
+                };
+                Type::List(Box::new(list_type))
+            }
+        }
+    }
+}
+
+impl ToString for RecordValue {
+    fn to_string(&self) -> String {
+        match self {
+            RecordValue::Int(i) => i.to_string(),
+            RecordValue::Float(f) => f.to_string(),
+            RecordValue::String(s) => s.clone(),
+            RecordValue::Boolean(b) => b.to_string(),
+            RecordValue::Null => String::new(),
+            RecordValue::Map(_) => String::new(),
+            RecordValue::List(_) => String::new(),
+        }
+    }
+}
+
+impl From<i64> for RecordValue {
+    fn from(num: i64) -> Self {
+        Self::Int(num)
+    }
+}
+
+impl From<f64> for RecordValue {
+    fn from(num: f64) -> Self {
+        Self::Float(num)
+    }
+}
+
+impl From<String> for RecordValue {
+    fn from(s: String) -> Self {
+        Self::String(s)
+    }
+}
+
+impl<'a> From<Cow<'a, str>> for RecordValue {
+    fn from(s: Cow<'a, str>) -> Self {
+        Self::String(s.into_owned())
+    }
+}
+
+impl From<&str> for RecordValue {
+    fn from(s: &str) -> Self {
+        Self::String(s.into())
+    }
+}
+
+impl From<bool> for RecordValue {
+    fn from(b: bool) -> Self {
+        Self::Boolean(b)
+    }
+}
+
+impl<T: Into<RecordValue>> From<Option<T>> for RecordValue {
+    fn from(opt: Option<T>) -> Self {
+        match opt {
+            Some(val) => val.into(),
+            None => Self::Null,
+        }
+    }
+}
+
+impl<T1: Into<RecordValue>, T2: Into<RecordValue>> From<HashMap<T1, T2>> for RecordValue {
+    fn from(map: HashMap<T1, T2>) -> Self {
+        let map = map.into_iter().map(|(k, v)| (k.into(), v.into())).collect();
+        Self::Map(map)
+    }
+}
+
+impl<T: Into<RecordValue>> From<Vec<T>> for RecordValue {
+    fn from(list: Vec<T>) -> Self {
+        let list = list.into_iter().map(|l| l.into()).collect();
+        Self::List(list)
+    }
+}

src/scripting/script.rs

@@ -1,9 +1,11 @@
-use serde::Serialize;
 use std::{marker::PhantomData, path::PathBuf};
 
 use crate::error::AppResult;
 
-use super::executor::{NuExecutor, VarValue};
+use super::{
+    executor::{NuExecutor, VarValue},
+    record::RecordValue,
+};
 
 /// A trait implemented for a given nu script type to
 /// associate arguments
@@ -19,7 +21,7 @@ pub trait Script {
 /// Script arguments that can be collected in a Vec to
 /// be passed to the script
 pub trait ScriptArgs {
-    fn get_args(self) -> Vec<String>;
+    fn get_args(self) -> Vec<RecordValue>;
 }
 
 /// A nu script instance that can be executed
@@ -49,17 +51,3 @@ impl<S: Script> NuScript<S> {
             .await
     }
 }
-
-pub struct JSONArgs<T: Serialize>(pub T);
-
-impl<T: Serialize> ScriptArgs for JSONArgs<T> {
-    fn get_args(self) -> Vec<String> {
-        // TODO: Make this lesss... weird
-        // Maybe try providing the value directly in the executor
-        // instead of parsing and wrapping it
-        vec![format!(
-            "('{}' | from json)",
-            serde_json::to_string(&self.0).unwrap()
-        )]
-    }
-}

src/tasks/setup_users.rs

@@ -1,6 +1,11 @@
+use std::collections::HashMap;
+
 use serde::Serialize;
 
-use crate::scripting::script::{JSONArgs, Script};
+use crate::scripting::{
+    record::RecordValue,
+    script::{Script, ScriptArgs},
+};
 
 pub struct SetupUsersScript;
 
@@ -18,9 +23,28 @@ pub struct User {
 }
 
 impl Script for SetupUsersScript {
-    type Args = JSONArgs<UsersConfig>;
+    type Args = UsersConfig;
 
     fn get_name() -> &'static str {
         "setup-users.nu"
     }
 }
+
+impl ScriptArgs for UsersConfig {
+    fn get_args(self) -> Vec<RecordValue> {
+        let mut user_cfg_map = HashMap::new();
+        let mut user_cfgs = Vec::new();
+
+        for user in self.users {
+            let mut user_map: HashMap<&'static str, RecordValue> = HashMap::new();
+            user_map.insert("name", user.name.into());
+            user_map.insert("password", user.password.into());
+            user_map.insert("sudoer", user.sudoer.into());
+            user_map.insert("shell", user.shell.into());
+            user_cfgs.push(user_map);
+        }
+        user_cfg_map.insert("users", user_cfgs);
+
+        vec![user_cfg_map.into()]
+    }
+}