fix: added completes to cmds

jdx · Dec 9, 2024 · f421d9e · f421d9e
1 parent 6783590
commit f421d9e
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 4 deletions.
diff --git a/cli/src/cli/complete_word.rs b/cli/src/cli/complete_word.rs
@@ -86,11 +86,11 @@ impl CompleteWord {
         } else if ctoken.starts_with('-') {
             self.complete_short_flag_names(&parsed.available_flags, &ctoken)
         } else if let Some(flag) = parsed.flag_awaiting_value {
-            self.complete_arg(&ctx, spec, flag.arg.as_ref().unwrap(), &ctoken)?
+            self.complete_arg(&ctx, spec, &parsed.cmd, flag.arg.as_ref().unwrap(), &ctoken)?
         } else {
             let mut choices = vec![];
             if let Some(arg) = parsed.cmd.args.get(parsed.args.len()) {
-                choices.extend(self.complete_arg(&ctx, spec, arg, &ctoken)?);
+                choices.extend(self.complete_arg(&ctx, spec, &parsed.cmd, arg, &ctoken)?);
             }
             if !parsed.cmd.subcommands.is_empty() {
                 choices.extend(self.complete_subcommands(&parsed.cmd, &ctoken));
@@ -184,14 +184,19 @@ impl CompleteWord {
         &self,
         ctx: &tera::Context,
         spec: &Spec,
+        cmd: &SpecCommand,
         arg: &SpecArg,
         ctoken: &str,
     ) -> miette::Result<Vec<(String, String)>> {
         static EMPTY_COMPL: Lazy<SpecComplete> = Lazy::new(SpecComplete::default);
 
         trace!("complete_arg: {arg} {ctoken}");
         let name = arg.name.to_lowercase();
-        let complete = spec.complete.get(&name).unwrap_or(&EMPTY_COMPL);
+        let complete = spec
+            .complete
+            .get(&name)
+            .or(cmd.complete.get(&name))
+            .unwrap_or(&EMPTY_COMPL);
         let type_ = complete.type_.as_ref().unwrap_or(&name);
 
         let builtin = self.complete_builtin(type_, ctoken);

diff --git a/lib/src/spec/cmd.rs b/lib/src/spec/cmd.rs
@@ -7,7 +7,7 @@ use crate::spec::context::ParsingContext;
 use crate::spec::helpers::NodeHelper;
 use crate::spec::is_false;
 use crate::spec::mount::SpecMount;
-use crate::{Spec, SpecArg, SpecFlag};
+use crate::{Spec, SpecArg, SpecComplete, SpecFlag};
 use indexmap::IndexMap;
 use itertools::Itertools;
 use kdl::{KdlDocument, KdlEntry, KdlNode, KdlValue};
@@ -48,6 +48,8 @@ pub struct SpecCommand {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub after_help_md: Option<String>,
     pub examples: Vec<SpecExample>,
+    #[serde(skip_serializing_if = "IndexMap::is_empty")]
+    pub complete: IndexMap<String, SpecComplete>,
 
     // TODO: make this non-public
     #[serde(skip)]
@@ -80,6 +82,7 @@ impl Default for SpecCommand {
             after_help_md: None,
             examples: vec![],
             subcommand_lookup: OnceLock::new(),
+            complete: IndexMap::new(),
         }
     }
 }
@@ -207,6 +210,10 @@ impl SpecCommand {
                         None => Some(child.arg(0)?.ensure_string()?),
                     }
                 }
+                "complete" => {
+                    let complete = SpecComplete::parse(ctx, &child)?;
+                    cmd.complete.insert(complete.name.clone(), complete);
+                }
                 k => bail_parse!(ctx, *child.node.name().span(), "unsupported cmd key {k}"),
             }
         }
@@ -314,6 +321,9 @@ impl SpecCommand {
         for (name, cmd) in other.subcommands {
             self.subcommands.insert(name, cmd);
         }
+        for (name, complete) in other.complete {
+            self.complete.insert(name, complete);
+        }
     }
 
     pub fn all_subcommands(&self) -> Vec<&SpecCommand> {
@@ -445,6 +455,10 @@ impl From<&SpecCommand> for KdlNode {
             let children = node.children_mut().get_or_insert_with(KdlDocument::new);
             children.nodes_mut().push(cmd.into());
         }
+        for complete in cmd.complete.values() {
+            let children = node.children_mut().get_or_insert_with(KdlDocument::new);
+            children.nodes_mut().push(complete.into());
+        }
         node
     }
 }

diff --git a/lib/src/spec/mod.rs b/lib/src/spec/mod.rs
@@ -346,6 +346,9 @@ impl Display for Spec {
         for complete in self.complete.values() {
             nodes.push(complete.into());
         }
+        for complete in self.cmd.complete.values() {
+            nodes.push(complete.into());
+        }
         for cmd in self.cmd.subcommands.values() {
             nodes.push(cmd.into())
         }