From fc82a3e52887b426302013a8bcd94fa76b9fcbf4 Mon Sep 17 00:00:00 2001
From: shannmu <shanmu1901@gmail.com>
Date: Fri, 16 Aug 2024 17:32:42 +0800
Subject: [PATCH] feat(complete): Support to complete custom external
 subcommand

---
 clap_complete/src/engine/complete.rs    | 30 ++++++++++++++--
 clap_complete/src/engine/custom.rs      | 47 +++++++++++++++++++++++++
 clap_complete/src/engine/mod.rs         |  1 +
 clap_complete/tests/testsuite/engine.rs | 12 +++++--
 4 files changed, 85 insertions(+), 5 deletions(-)

diff --git a/clap_complete/src/engine/complete.rs b/clap_complete/src/engine/complete.rs
index 2c5de4f8170..c5ff943dbc5 100644
--- a/clap_complete/src/engine/complete.rs
+++ b/clap_complete/src/engine/complete.rs
@@ -7,6 +7,7 @@ use super::custom::complete_path;
 use super::ArgValueCandidates;
 use super::ArgValueCompleter;
 use super::CompletionCandidate;
+use super::SubcommandCandidates;
 
 /// Complete the given command, shell-agnostic
 pub fn complete(
@@ -414,10 +415,35 @@ fn complete_subcommand(value: &str, cmd: &clap::Command) -> Vec<CompletionCandid
         value
     );
 
-    subcommands(cmd)
+    let mut scs: Vec<CompletionCandidate> = subcommands(cmd)
         .into_iter()
         .filter(|x| x.get_value().starts_with(value))
-        .collect()
+        .collect();
+    if cmd.is_allow_external_subcommands_set() {
+        let external_completer = cmd.get::<SubcommandCandidates>();
+        if let Some(completer) = external_completer {
+            scs.extend(complete_external_subcommand(value, completer));
+        }
+    }
+
+    scs.sort();
+    scs.dedup();
+    scs
+}
+
+fn complete_external_subcommand(
+    value: &str,
+    completer: &SubcommandCandidates,
+) -> Vec<CompletionCandidate> {
+    debug!("complete_custom_arg_value: completer={completer:?}, value={value:?}");
+
+    let mut values = Vec::new();
+    let custom_arg_values = completer.candidates();
+    values.extend(custom_arg_values);
+
+    values.retain(|comp| comp.get_value().starts_with(value));
+
+    values
 }
 
 /// Gets all the long options, their visible aliases and flags of a [`clap::Command`] with formatted `--` prefix.
diff --git a/clap_complete/src/engine/custom.rs b/clap_complete/src/engine/custom.rs
index 00dde78d67c..db0eb2a17c9 100644
--- a/clap_complete/src/engine/custom.rs
+++ b/clap_complete/src/engine/custom.rs
@@ -3,6 +3,7 @@ use std::ffi::OsStr;
 use std::sync::Arc;
 
 use clap::builder::ArgExt;
+use clap::builder::CommandExt;
 use clap_lex::OsStrExt as _;
 
 use super::CompletionCandidate;
@@ -131,8 +132,54 @@ impl std::fmt::Debug for ArgValueCandidates {
 
 impl ArgExt for ArgValueCandidates {}
 
+/// Extend [`Command`][clap::Command] with a [`ValueCandidates`]
+///
+/// # Example
+/// ```rust
+/// use clap::Parser;
+/// use clap_complete::engine::{SubcommandCandidates, CompletionCandidate};
+/// #[derive(Debug, Parser)]
+/// #[clap(name = "cli", add = SubcommandCandidates::new(|| { vec![
+///     CompletionCandidate::new("foo"),
+///     CompletionCandidate::new("bar"),
+///     CompletionCandidate::new("baz")] }))]
+/// struct Cli {
+///     #[arg(long)]
+///     input: Option<String>,
+/// }
+/// ```
+#[derive(Clone)]
+pub struct SubcommandCandidates(Arc<dyn ValueCandidates>);
+
+impl SubcommandCandidates {
+    /// Create a new `SubcommandCandidates` with a custom completer
+    pub fn new<C>(completer: C) -> Self
+    where
+        C: ValueCandidates + 'static,
+    {
+        Self(Arc::new(completer))
+    }
+
+    /// All potential candidates for an external subcommand.
+    ///
+    /// See [`CompletionCandidate`] for more information.
+    pub fn candidates(&self) -> Vec<CompletionCandidate> {
+        self.0.candidates()
+    }
+}
+
+impl std::fmt::Debug for SubcommandCandidates {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_str(type_name::<Self>())
+    }
+}
+
+impl CommandExt for SubcommandCandidates {}
+
 /// User-provided completion candidates for an [`Arg`][clap::Arg], see [`ArgValueCandidates`]
 ///
+/// User-provided completion candidates for an [`Subcommand`][clap::Subcommand], see [`SubcommandCandidates`]
+///
 /// This is useful when predefined value hints are not enough.
 pub trait ValueCandidates: Send + Sync {
     /// All potential candidates for an argument.
diff --git a/clap_complete/src/engine/mod.rs b/clap_complete/src/engine/mod.rs
index 43a1f9f0c10..358f0653e53 100644
--- a/clap_complete/src/engine/mod.rs
+++ b/clap_complete/src/engine/mod.rs
@@ -11,5 +11,6 @@ pub use complete::complete;
 pub use custom::ArgValueCandidates;
 pub use custom::ArgValueCompleter;
 pub use custom::PathCompleter;
+pub use custom::SubcommandCandidates;
 pub use custom::ValueCandidates;
 pub use custom::ValueCompleter;
diff --git a/clap_complete/tests/testsuite/engine.rs b/clap_complete/tests/testsuite/engine.rs
index 1848d58bf5c..a3a18912a4e 100644
--- a/clap_complete/tests/testsuite/engine.rs
+++ b/clap_complete/tests/testsuite/engine.rs
@@ -5,7 +5,7 @@ use std::path::Path;
 
 use clap::{builder::PossibleValue, Command};
 use clap_complete::engine::{
-    ArgValueCandidates, ArgValueCompleter, CompletionCandidate, PathCompleter,
+    ArgValueCandidates, ArgValueCompleter, CompletionCandidate, PathCompleter, SubcommandCandidates,
 };
 use snapbox::assert_data_eq;
 
@@ -1080,19 +1080,25 @@ pos_b
 #[test]
 fn suggest_external_subcommand() {
     let mut cmd = Command::new("dynamic")
+        .allow_external_subcommands(true)
+        .add(SubcommandCandidates::new(|| {
+            vec![CompletionCandidate::new("external")]
+        }))
         .arg(clap::Arg::new("positional").value_parser(["pos1", "pos2", "pos3"]));
 
     assert_data_eq!(
         complete!(cmd, " [TAB]"),
         snapbox::str![
-            "--help\tPrint help
--h\tPrint help
+            "external
 pos1
 pos2
 pos3
+--help\tPrint help
 "
         ]
     );
+
+    assert_data_eq!(complete!(cmd, "e[TAB]"), snapbox::str!["external"]);
 }
 
 #[test]