feat: support symbol parameter (#260)

src/command/mod.rs

@@ -7,11 +7,11 @@ use self::root_data::RootData;
 use crate::argc_value::ArgcValue;
 use crate::matcher::Matcher;
 use crate::param::{FlagOptionParam, PositionalParam};
-use crate::parser::{parse, Event, EventData, EventScope, Position};
+use crate::parser::{parse, parse_symbol, Event, EventData, EventScope, Position};
 use crate::utils::INTERNAL_MODE;
 use crate::Result;
 
-use anyhow::bail;
+use anyhow::{anyhow, bail};
 use indexmap::IndexMap;
 use std::cell::RefCell;
 use std::collections::HashMap;
@@ -48,14 +48,15 @@ pub struct Command {
     pub(crate) names_checker: NamesChecker,
     pub(crate) root: Arc<RefCell<RootData>>,
     pub(crate) aliases: Vec<String>,
-    pub(crate) metadata: IndexMap<String, (String, Position)>,
+    pub(crate) metadata: Vec<(String, String, Position)>,
+    pub(crate) symbols: IndexMap<char, SymbolParam>,
 }
 
 impl Command {
     pub fn new(source: &str) -> Result<Self> {
         let events = parse(source)?;
         let mut root = Command::new_from_events(&events)?;
-        if root.metadata.contains_key("inherit-flag-options") {
+        if root.has_metadata("inherit-flag-options") {
             root.inherit_flag_options();
         }
         Ok(root)
@@ -102,11 +103,8 @@ impl Command {
             .collect();
         let positional_params: Vec<serde_json::Value> =
             self.positional_params.iter().map(|v| v.to_json()).collect();
-        let mut metadata = serde_json::Map::new();
-        for (k, (v, _)) in &self.metadata {
-            metadata.insert(k.to_string(), serde_json::Value::String(v.to_string()));
-        }
-        let metadata = serde_json::Value::Object(metadata);
+        let metadata: Vec<Vec<&String>> =
+            self.metadata.iter().map(|(k, v, _)| vec![k, v]).collect();
         serde_json::json!({
             "describe": self.describe,
             "name": self.name,
@@ -140,15 +138,14 @@ impl Command {
                 }
                 EventData::Meta(key, value) => {
                     let cmd = Self::get_cmd(&mut root_cmd, "@meta", position)?;
-                    if let Some((_, pos)) = cmd.metadata.get(&key) {
-                        bail!(
-                            "@meta(line {}) conflicts with '{}' at line {}",
-                            position,
-                            key,
-                            pos
-                        )
+                    if key == "symbol" {
+                        let (ch, name, choice_fn) = parse_symbol(&value).ok_or_else(|| {
+                            anyhow!("@meta(line {}) invalid symbol value", position)
+                        })?;
+                        cmd.symbols
+                            .insert(ch, (name.to_string(), choice_fn.map(|v| v.to_string())));
                     }
-                    cmd.metadata.insert(key, (value, position));
+                    cmd.metadata.push((key, value, position));
                 }
                 EventData::Cmd(value) => {
                     if root_data.borrow().scope == EventScope::CmdStart {
@@ -270,6 +267,10 @@ impl Command {
         Ok(root_cmd)
     }
 
+    pub(crate) fn has_metadata(&self, key: &str) -> bool {
+        self.metadata.iter().any(|(k, _, _)| k == key)
+    }
+
     pub(crate) fn render_help(&self, cmd_paths: &[&str], term_width: Option<usize>) -> String {
         let mut output = vec![];
         if self.version.is_some() {
@@ -644,6 +645,8 @@ impl Command {
     }
 }
 
+pub(crate) type SymbolParam = (String, Option<String>);
+
 fn retrive_cmd<'a>(cmd: &'a mut Command, cmd_paths: &[String]) -> Option<&'a mut Command> {
     if cmd_paths.is_empty() {
         return Some(cmd);

src/matcher.rs

@@ -5,7 +5,7 @@ use std::{
 };
 
 use crate::{
-    command::Command,
+    command::{Command, SymbolParam},
     compgen::CompColor,
     param::{ChoiceData, FlagOptionParam, ParamData, PositionalParam},
     utils::run_param_fns,
@@ -22,6 +22,7 @@ pub(crate) struct Matcher<'a, 'b> {
     args: &'b [String],
     flag_option_args: Vec<Vec<FlagOptionArg<'a, 'b>>>,
     positional_args: Vec<&'b str>,
+    symbol_args: Vec<SymbolArg<'a, 'b>>,
     dashes: Option<usize>,
     arg_comp: ArgComp,
     choice_fns: HashSet<&'a str>,
@@ -32,6 +33,7 @@ pub(crate) struct Matcher<'a, 'b> {
 }
 
 type FlagOptionArg<'a, 'b> = (&'b str, Vec<&'b str>, Option<&'a str>); // key, values, param_name
+type SymbolArg<'a, 'b> = (&'b str, &'a SymbolParam);
 type LevelCommand<'a, 'b> = (&'b str, &'a Command, usize); // name, command, arg_index
 
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -40,6 +42,7 @@ pub(crate) enum ArgComp {
     FlagOrOptionCombine(String),
     CommandOrPositional,
     OptionValue(String, usize),
+    Symbol(char),
     Any,
 }
 
@@ -65,12 +68,13 @@ pub(crate) enum MatchError {
 
 impl<'a, 'b> Matcher<'a, 'b> {
     pub(crate) fn new(root: &'a Command, args: &'b [String], compgen: bool) -> Self {
-        let combine_shorts = root.metadata.contains_key("combine-shorts");
+        let combine_shorts = root.has_metadata("combine-shorts");
         let mut cmds: Vec<LevelCommand> = vec![(args[0].as_str(), root, 0)];
         let mut cmd_level = 0;
         let mut arg_index = 1;
         let mut flag_option_args = vec![vec![]];
         let mut positional_args = vec![];
+        let mut symbol_args = vec![];
         let mut dashes = None;
         let mut split_last_arg_at = None;
         let mut arg_comp = ArgComp::Any;
@@ -243,6 +247,15 @@ impl<'a, 'b> Matcher<'a, 'b> {
                         arg_index,
                         &mut is_rest_args_positional,
                     );
+                } else if let Some((ch, symbol_param)) = find_symbol(cmd, arg) {
+                    if let Some(choice_fn) = &symbol_param.1 {
+                        choice_fns.insert(choice_fn);
+                    }
+                    symbol_args.push((&arg[1..], symbol_param));
+                    if is_last_arg {
+                        arg_comp = ArgComp::Symbol(ch);
+                        split_last_arg_at = Some(1);
+                    }
                 } else {
                     add_positional_arg(
                         &mut positional_args,
@@ -273,6 +286,7 @@ impl<'a, 'b> Matcher<'a, 'b> {
             args,
             flag_option_args,
             positional_args,
+            symbol_args,
             dashes,
             arg_comp,
             choice_fns,
@@ -405,6 +419,7 @@ impl<'a, 'b> Matcher<'a, 'b> {
                     vec![]
                 }
             }
+            ArgComp::Symbol(ch) => comp_symbol(last_cmd, *ch),
             ArgComp::Any => {
                 if self.positional_args.len() == 2 && self.positional_args[0] == "help" {
                     return comp_subcomands(last_cmd, false);
@@ -437,6 +452,11 @@ impl<'a, 'b> Matcher<'a, 'b> {
         let level = cmds_len - 1;
         let last_cmd = self.cmds[level].1;
         let cmd_arg_index = self.cmds[level].2;
+
+        for (arg, (name, _)) in self.symbol_args.iter() {
+            output.push(ArgcValue::Single(name.to_string(), arg.to_string()));
+        }
+
         for level in 0..cmds_len {
             let args = &self.flag_option_args[level];
             let cmd = self.cmds[level].1;
@@ -866,6 +886,7 @@ impl<'a, 'b> Matcher<'a, 'b> {
     }
 }
 
+/// (value, description, nospace, color)
 pub(crate) type CompItem = (String, String, bool, CompColor);
 
 fn find_subcommand<'a>(
@@ -882,6 +903,15 @@ fn find_subcommand<'a>(
     })
 }
 
+fn find_symbol<'a>(cmd: &'a Command, arg: &str) -> Option<(char, &'a SymbolParam)> {
+    for (ch, param) in cmd.symbols.iter() {
+        if arg.starts_with(*ch) {
+            return Some((*ch, param));
+        }
+    }
+    None
+}
+
 fn add_positional_arg<'a>(
     positional_args: &mut Vec<&'a str>,
     arg: &'a str,
@@ -1090,6 +1120,31 @@ fn comp_subcomands(cmd: &Command, flag: bool) -> Vec<CompItem> {
     output
 }
 
+fn comp_symbol(cmd: &Command, ch: char) -> Vec<CompItem> {
+    if let Some((name, choices_fn)) = cmd.symbols.get(&ch) {
+        match choices_fn {
+            Some(choices_fn) => {
+                vec![(
+                    format!("__argc_fn={}", choices_fn),
+                    String::new(),
+                    false,
+                    CompColor::of_value(),
+                )]
+            }
+            None => {
+                vec![(
+                    format!("__argc_value={}", name),
+                    String::new(),
+                    false,
+                    CompColor::of_value(),
+                )]
+            }
+        }
+    } else {
+        vec![]
+    }
+}
+
 fn comp_flag_option(param: &FlagOptionParam, index: usize) -> Vec<CompItem> {
     let value_name = param
         .arg_value_names

src/parser.rs

@@ -109,6 +109,11 @@ pub(crate) fn parse(source: &str) -> Result<Vec<Event>> {
     Ok(result)
 }
 
+pub(crate) fn parse_symbol(input: &str) -> Option<(char, &str, Option<&str>)> {
+    let input = input.trim();
+    parse_symbol_data(input).map(|(_, v)| v).ok()
+}
+
 fn parse_line(line: &str) -> nom::IResult<&str, Option<Option<EventData>>> {
     alt((map(alt((parse_tag, parse_fn)), Some), success(None)))(line)
 }
@@ -621,6 +626,20 @@ fn parse_normal_comment(input: &str) -> nom::IResult<&str, &str> {
     ))(input)
 }
 
+fn parse_symbol_data(input: &str) -> nom::IResult<&str, (char, &str, Option<&str>)> {
+    map(
+        terminated(
+            tuple((
+                alt((char('@'), char('+'))),
+                parse_name,
+                opt(delimited(char('['), parse_value_fn, char(']'))),
+            )),
+            eof,
+        ),
+        |(symbol, name, choice_fn)| (symbol, name, choice_fn),
+    )(input)
+}
+
 fn notation_text(input: &str, balances: usize) -> nom::IResult<&str, usize> {
     let (i1, c1) = anychar(input)?;
     match c1 {
@@ -971,4 +990,16 @@ mod tests {
         assert_token!("foo=bar", Ignore);
         assert_token!("#!/bin/bash", Ignore);
     }
+
+    #[test]
+    fn test_parse_symbol() {
+        assert_eq!(
+            parse_symbol("+toolchain").unwrap(),
+            ('+', "toolchain", None)
+        );
+        assert_eq!(
+            parse_symbol("+toolchain[`_choice_toolchain`]").unwrap(),
+            ('+', "toolchain", Some("_choice_toolchain"))
+        );
+    }
 }

tests/compgen.rs

@@ -79,6 +79,25 @@ fn shorts() {
     );
 }
 
+#[test]
+fn symbol() {
+    let script = r###"
+# @meta symbol +toolchain[`_choice_fn`]
+# @meta symbol @file
+# @option --oa
+
+_choice_fn() {
+    echo stable
+    echo beta
+    echo nightly
+}
+"###;
+    snapshot_compgen!(
+        script,
+        [vec!["prog", "+"], vec!["prog", "@"], vec!["prog", "+s"]]
+    );
+}
+
 #[test]
 fn subcmds() {
     const SCRIPT: &str = r###"

tests/snapshots/integration__compgen__symbol.snap

@@ -0,0 +1,16 @@
+---
+source: tests/compgen.rs
+expression: data
+---
+************ COMPGEN `prog +` ************
+stable	/color:default
+beta	/color:default
+nightly	/color:default
+
+************ COMPGEN `prog @` ************
+__argc_value=file	/color:default
+
+************ COMPGEN `prog +s` ************
+stable	/color:default
+
+

tests/snapshots/integration__export__case1.snap

@@ -7,9 +7,12 @@ expression: output
   "name": null,
   "author": null,
   "version": null,
-  "metadata": {
-    "combine-shorts": ""
-  },
+  "metadata": [
+    [
+      "combine-shorts",
+      ""
+    ]
+  ],
   "options": [],
   "positionals": [],
   "aliases": [],
@@ -19,7 +22,7 @@ expression: output
       "name": "cmda",
       "author": null,
       "version": null,
-      "metadata": {},
+      "metadata": [],
       "options": [
         {
           "name": "a",
@@ -268,7 +271,7 @@ expression: output
       "name": "cmdb",
       "author": null,
       "version": null,
-      "metadata": {},
+      "metadata": [],
       "options": [
         {
           "name": "oa",
@@ -367,7 +370,7 @@ expression: output
       "name": "cmdc",
       "author": null,
       "version": null,
-      "metadata": {},
+      "metadata": [],
       "options": [
         {
           "name": "oe",
@@ -425,7 +428,7 @@ expression: output
       "name": "cmdd",
       "author": null,
       "version": null,
-      "metadata": {},
+      "metadata": [],
       "options": [
         {
           "name": "oa",
@@ -451,7 +454,7 @@ expression: output
       "name": "cmde",
       "author": null,
       "version": null,
-      "metadata": {},
+      "metadata": [],
       "options": [
         {
           "name": "oa",