feat: validate param with choices fn (#100)

* feat: validate param with choices fn

* fix tests

* cargo fmt

.github/workflows/ci.yaml

@@ -40,6 +40,8 @@ jobs:
 
     - uses: Swatinem/rust-cache@v1
 
+    - run: cargo install --path .
+
     - name: Test
       run: cargo test --all
 

src/bin/argc/main.rs

@@ -39,7 +39,7 @@ fn run() -> Result<i32> {
         match argc_cmd {
             "--argc-eval" => {
                 let (source, cmd_args) = parse_script_args(&args[2..])?;
-                let values = argc::eval(&source, &cmd_args)?;
+                let values = argc::eval(Some(&args[2]), &source, &cmd_args)?;
                 println!("{}", argc::ArgcValue::to_shell(values))
             }
             "--argc-create" => {

src/command/mod.rs

@@ -17,9 +17,13 @@ use std::collections::HashMap;
 use std::result::Result as StdResult;
 use std::sync::Arc;
 
-pub fn eval(source: &str, args: &[String]) -> Result<Vec<ArgcValue>> {
-    let mut cmd = Command::new(source)?;
-    cmd.eval(args)
+pub fn eval(
+    script_path: Option<&str>,
+    script_content: &str,
+    args: &[String],
+) -> Result<Vec<ArgcValue>> {
+    let mut cmd = Command::new(script_content)?;
+    cmd.eval(script_path, args)
 }
 
 pub fn export(source: &str, name: &str) -> Result<serde_json::Value> {
@@ -52,7 +56,7 @@ impl Command {
         Command::new_from_events(&events)
     }
 
-    pub fn eval(&mut self, args: &[String]) -> Result<Vec<ArgcValue>> {
+    pub fn eval(&mut self, script_path: Option<&str>, args: &[String]) -> Result<Vec<ArgcValue>> {
         if args.is_empty() {
             bail!("Invalid args");
         }
@@ -78,7 +82,10 @@ impl Command {
             arg_values.push(ArgcValue::ParamFn(args[1].clone()));
             return Ok(arg_values);
         }
-        let matcher = Matcher::new(self, args);
+        let mut matcher = Matcher::new(self, args);
+        if let Some(script_path) = script_path {
+            matcher.set_script_path(script_path)
+        }
         Ok(matcher.to_arg_values())
     }
 

src/compgen.rs

@@ -1,9 +1,9 @@
 use crate::command::Command;
 use crate::matcher::Matcher;
-use crate::utils::{escape_shell_words, get_shell_path, split_shell_words};
+use crate::utils::{escape_shell_words, run_param_fns, split_shell_words};
 use crate::Result;
 use anyhow::bail;
-use std::{process, str::FromStr};
+use std::str::FromStr;
 
 pub fn compgen(
     shell: Shell,
@@ -168,32 +168,25 @@ fn expand_candicates(
         }
     }
     if !param_fns.is_empty() {
-        if let Some(shell) = get_shell_path() {
-            for param_fn in param_fns {
-                if let Ok(fn_output) = process::Command::new(&shell)
-                    .arg(script_file)
-                    .arg(&param_fn)
-                    .arg(line)
-                    .output()
-                {
-                    let fn_output = String::from_utf8_lossy(&fn_output.stdout);
-                    for fn_output_line in fn_output.split('\n') {
-                        let output_line = fn_output_line.trim();
-                        if !output_line.is_empty()
-                            && (output_line.starts_with("__argc_")
-                                || output_line.starts_with(filter))
-                        {
-                            if let Some((x, y)) = output_line.split_once('\t') {
-                                output.push((x.to_string(), y.to_string()));
-                            } else {
-                                output.push((output_line.to_string(), String::new()));
-                            }
+        let fns: Vec<&str> = param_fns.iter().map(|v| v.as_str()).collect();
+        if let Some(param_fn_outputs) = run_param_fns(script_file, &fns, line) {
+            for param_fn_output in param_fn_outputs {
+                for output_line in param_fn_output.split('\n') {
+                    let output_line = output_line.trim();
+                    if !output_line.is_empty()
+                        && (output_line.starts_with("__argc_") || output_line.starts_with(filter))
+                    {
+                        if let Some((x, y)) = output_line.split_once('\t') {
+                            output.push((x.to_string(), y.to_string()));
+                        } else {
+                            output.push((output_line.to_string(), String::new()));
                         }
                     }
                 }
             }
         }
     }
+
     if output.len() == 1 {
         let value = &output[0].0;
         if let Some(value_name) = value.strip_prefix("__argc_value") {

src/matcher.rs

@@ -1,8 +1,9 @@
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 
 use crate::{
     command::Command,
     param::{FlagOptionParam, PositionalParam},
+    utils::run_param_fns,
     ArgcValue,
 };
 
@@ -15,6 +16,9 @@ pub struct Matcher<'a, 'b> {
     positional_args: Vec<&'b str>,
     dashdash: Option<usize>,
     arg_comp: ArgComp,
+    choices_fns: HashSet<&'a str>,
+    choices_values: HashMap<&'a str, Vec<String>>,
+    script_path: Option<String>,
 }
 
 type FlagOptionArg<'a, 'b> = (&'b str, Vec<&'b str>, Option<&'a str>);
@@ -51,6 +55,7 @@ impl<'a, 'b> Matcher<'a, 'b> {
         let mut positional_args = vec![];
         let mut dashdash = None;
         let mut arg_comp = ArgComp::Any;
+        let mut choices_fns = HashSet::new();
         let args_len = args.len();
         if let Some(arg) = args.last() {
             if arg.starts_with('-') {
@@ -77,8 +82,14 @@ impl<'a, 'b> Matcher<'a, 'b> {
                             arg_comp = ArgComp::OptionValue(param.name.clone(), 0)
                         }
                     }
+                    if let Some(choices_fn) = param.and_then(|v| v.choices_fn.as_ref()) {
+                        choices_fns.insert(choices_fn.as_str());
+                    }
                     flag_option_args[cmd_level].push((k, vec![v], param.map(|v| v.name.as_str())));
                 } else if let Some(param) = cmd.find_flag_option(arg) {
+                    if let Some(choices_fn) = param.choices_fn.as_ref() {
+                        choices_fns.insert(choices_fn.as_str());
+                    }
                     match_flag_option(
                         &mut flag_option_args[cmd_level],
                         args,
@@ -89,6 +100,9 @@ impl<'a, 'b> Matcher<'a, 'b> {
                 } else if let Some(mut list) = match_combine_shorts(cmd, arg) {
                     let name = list.pop().and_then(|v| v.2).unwrap();
                     let param = cmd.find_flag_option(name).unwrap();
+                    if let Some(choices_fn) = param.choices_fn.as_ref() {
+                        choices_fns.insert(choices_fn.as_str());
+                    }
                     flag_option_args[cmd_level].extend(list);
                     match_flag_option(
                         &mut flag_option_args[cmd_level],
@@ -113,12 +127,39 @@ impl<'a, 'b> Matcher<'a, 'b> {
             }
             arg_index += 1;
         }
+        let last_cmd = cmds.last().unwrap().1;
+        choices_fns.extend(
+            last_cmd
+                .positional_params
+                .iter()
+                .filter_map(|v| v.choices_fn.as_deref()),
+        );
         Self {
             cmds,
             flag_option_args,
             positional_args,
             dashdash,
             arg_comp,
+            choices_fns,
+            choices_values: HashMap::new(),
+            script_path: None,
+        }
+    }
+
+    pub fn set_script_path(&mut self, script_path: &str) {
+        self.script_path = Some(script_path.to_string());
+        let fns: Vec<&str> = self.choices_fns.iter().copied().collect();
+        if let Some(list) = run_param_fns(script_path, &fns, "") {
+            for (i, fn_output) in list.into_iter().enumerate() {
+                let choices: Vec<String> = fn_output
+                    .split('\n')
+                    .map(|v| v.trim().to_string())
+                    .filter(|v| !v.is_empty())
+                    .collect();
+                if !choices.is_empty() {
+                    self.choices_values.insert(fns[i], choices);
+                }
+            }
         }
     }
 
@@ -293,6 +334,7 @@ impl<'a, 'b> Matcher<'a, 'b> {
                 return Some(MatchError::InvalidSubcommand);
             }
         }
+
         let positional_values = self.match_positionals();
         let positional_values_len = positional_values.len();
         let positional_params_len = last_cmd.positional_params.len();
@@ -315,7 +357,7 @@ impl<'a, 'b> Matcher<'a, 'b> {
         for (i, param) in last_cmd.positional_params.iter().enumerate() {
             if let (Some(value), Some(choices)) = (
                 positional_values.get(i).and_then(|v| v.first()),
-                &param.choices,
+                get_param_choices(&param.choices, &param.choices_fn, &self.choices_values),
             ) {
                 if !choices.contains(&value.to_string()) {
                     return Some(MatchError::InvalidValue(
@@ -378,7 +420,11 @@ impl<'a, 'b> Matcher<'a, 'b> {
                                 ));
                             }
                         }
-                        if let Some(choices) = &param.choices {
+                        if let Some(choices) = get_param_choices(
+                            &param.choices,
+                            &param.choices_fn,
+                            &self.choices_values,
+                        ) {
                             let value = values[0];
                             if !choices.contains(&value.to_string()) {
                                 return Some(MatchError::InvalidValue(
@@ -728,3 +774,15 @@ fn comp_param(
         vec![(value, describe.into())]
     }
 }
+
+fn get_param_choices<'a, 'b: 'a>(
+    choices: &'a Option<Vec<String>>,
+    choices_fn: &'a Option<String>,
+    choices_values: &'a HashMap<&str, Vec<String>>,
+) -> Option<&'a Vec<String>> {
+    choices.as_ref().or_else(|| {
+        choices_fn
+            .as_ref()
+            .and_then(|fn_name| choices_values.get(fn_name.as_str()))
+    })
+}

src/utils.rs

@@ -3,6 +3,7 @@ use convert_case::{Boundary, Converter, Pattern};
 use std::{
     env,
     path::{Path, PathBuf},
+    process, thread,
 };
 use which::which;
 
@@ -66,6 +67,33 @@ pub fn get_bash_path() -> Option<PathBuf> {
     which("bash").ok()
 }
 
+pub fn run_param_fns(script_file: &str, param_fns: &[&str], line: &str) -> Option<Vec<String>> {
+    let shell = get_shell_path()?;
+    let handles: Vec<_> = param_fns
+        .iter()
+        .map(|param_fn| {
+            let script_file = script_file.to_string();
+            let line = line.to_string();
+            let param_fn = param_fn.to_string();
+            let shell = shell.clone();
+            thread::spawn(move || {
+                process::Command::new(shell)
+                    .arg(script_file)
+                    .arg(param_fn)
+                    .arg(line)
+                    .output()
+                    .ok()
+                    .map(|out| String::from_utf8_lossy(&out.stdout).to_string())
+            })
+        })
+        .collect();
+    let list: Vec<String> = handles
+        .into_iter()
+        .map(|h| h.join().ok().flatten().unwrap_or_default())
+        .collect();
+    Some(list)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

tests/compgen.rs

@@ -132,3 +132,8 @@ fn test_compgen_cmd_single_dash() {
 fn test_compgen_cmd_single_dash2() {
     snapshot_compgen!("cmd_single_dash -f");
 }
+
+#[test]
+fn test_compgen_choice_fns() {
+    snapshot_compgen!("cmd_option_names --op10 ");
+}

tests/macros.rs

@@ -1,11 +1,12 @@
 #[macro_export]
 macro_rules! snapshot {
     (
+		$path:expr,
         $source:expr,
         $args:expr
     ) => {
         let args: Vec<String> = $args.iter().map(|v| v.to_string()).collect();
-        let values = argc::eval($source, &args).unwrap();
+        let values = argc::eval($path, $source, &args).unwrap();
         let output = argc::ArgcValue::to_shell(values);
         let args = $args.join(" ");
         let output = format!(
@@ -21,6 +22,14 @@ OUTPUT
     };
 }
 
+#[macro_export]
+macro_rules! snapshot_spec {
+    ($args:expr) => {
+        let (path, source) = $crate::fixtures::get_spec();
+        snapshot!(Some(path.as_str()), source.as_str(), $args);
+    };
+}
+
 #[macro_export]
 macro_rules! plain {
     (
@@ -29,7 +38,7 @@ macro_rules! plain {
 		$output:expr
     ) => {
         let args: Vec<String> = $args.iter().map(|v| v.to_string()).collect();
-        let values = argc::eval($source, &args).unwrap();
+        let values = argc::eval(None, $source, &args).unwrap();
         let output = argc::ArgcValue::to_shell(values);
         assert_eq!(output, $output);
     };
@@ -43,7 +52,7 @@ macro_rules! fatal {
         $err:expr
     ) => {
         let args: Vec<String> = $args.iter().map(|v| v.to_string()).collect();
-        let err = argc::eval($source, &args).unwrap_err();
+        let err = argc::eval(None, $source, &args).unwrap_err();
         assert_eq!(err.to_string().as_str(), $err);
     };
 }

tests/main_fn_test.rs

@@ -33,7 +33,7 @@ main() {
 }
     "###;
     plain!(script, &["prog", "cmd"], "cmd");
-    snapshot!(script, &["prog", "-h"]);
+    snapshot!(None, script, &["prog", "-h"]);
 }
 
 #[test]
@@ -48,7 +48,7 @@ cmd() {
 
     "###;
     plain!(script, &["prog", "cmd"], "cmd");
-    snapshot!(script, &["prog"]);
+    snapshot!(None, script, &["prog"]);
 }
 
 #[test]
@@ -62,5 +62,5 @@ cmd() {
 }
 
     "###;
-    snapshot!(script, &["prog", "-h"]);
+    snapshot!(None, script, &["prog", "-h"]);
 }