From 3239c02cba73453a207f5aa5fe168f75bac5a619 Mon Sep 17 00:00:00 2001
From: shanmu <shanmu1901@gmail.com>
Date: Fri, 2 Aug 2024 23:18:47 +0800
Subject: [PATCH] feat(clap_complete): Support to complete custom value of
 argument

---
 Cargo.toml                               |  2 +-
 clap_complete/Cargo.toml                 |  2 +-
 clap_complete/src/dynamic/completer.rs   | 80 +++++++++++++++++++++++-
 clap_complete/tests/testsuite/dynamic.rs | 29 ++++++++-
 4 files changed, 108 insertions(+), 5 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index f4f3f97f134..d8706089b66 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -168,7 +168,7 @@ string = ["clap_builder/string"]  # Allow runtime generated strings
 
 # In-work features
 unstable-v5 = ["clap_builder/unstable-v5", "clap_derive?/unstable-v5", "deprecated"]
-unstable-ext = []
+unstable-ext = ["clap_builder/unstable-ext"]
 unstable-styles = ["clap_builder/unstable-styles"]  # deprecated
 
 [lib]
diff --git a/clap_complete/Cargo.toml b/clap_complete/Cargo.toml
index 495ba99ab98..c24e0749462 100644
--- a/clap_complete/Cargo.toml
+++ b/clap_complete/Cargo.toml
@@ -57,7 +57,7 @@ required-features = ["unstable-dynamic"]
 [features]
 default = []
 unstable-doc = ["unstable-dynamic"] # for docs.rs
-unstable-dynamic = ["dep:clap_lex", "dep:shlex", "dep:unicode-xid", "clap/derive", "dep:is_executable", "dep:pathdiff"]
+unstable-dynamic = ["dep:clap_lex", "dep:shlex", "dep:unicode-xid", "clap/derive", "dep:is_executable", "dep:pathdiff", "clap/unstable-ext"]
 debug = ["clap/debug"]
 
 [lints]
diff --git a/clap_complete/src/dynamic/completer.rs b/clap_complete/src/dynamic/completer.rs
index 44837d758df..e02ffa6718b 100644
--- a/clap_complete/src/dynamic/completer.rs
+++ b/clap_complete/src/dynamic/completer.rs
@@ -1,7 +1,9 @@
-use core::num;
+use std::any::type_name;
 use std::ffi::OsStr;
 use std::ffi::OsString;
+use std::sync::Arc;
 
+use clap::builder::ArgExt;
 use clap::builder::StyledStr;
 use clap_lex::OsStrExt as _;
 
@@ -385,6 +387,10 @@ fn complete_arg_value(
                 values.extend(complete_path(value_os, current_dir, |_| true));
             }
         }
+
+        // Add custom completion at the same level as the ValueHint.
+        values.extend(complete_custom_arg_value(value_os, arg));
+
         values.sort();
     }
 
@@ -442,6 +448,20 @@ fn complete_path(
     completions
 }
 
+fn complete_custom_arg_value(value: &OsStr, arg: &clap::Arg) -> Vec<CompletionCandidate> {
+    let mut values = Vec::new();
+    debug!("complete_custom_arg_value: arg={arg:?}, value={value:?}");
+
+    if let Some(completer) = arg.get::<ArgValueCompleter>() {
+        let custom_arg_values = completer.0.completions();
+        values.extend(custom_arg_values);
+    }
+
+    values.retain(|comp| comp.get_content().starts_with(&value.to_string_lossy()));
+
+    values
+}
+
 fn complete_subcommand(value: &str, cmd: &clap::Command) -> Vec<CompletionCandidate> {
     debug!(
         "complete_subcommand: cmd={:?}, value={:?}",
@@ -701,3 +721,61 @@ impl CompletionCandidate {
         self.hidden
     }
 }
+
+/// User-provided completion candidates for an argument.
+///
+/// This is useful when predefined value hints are not enough.
+pub trait CustomCompleter: Send + Sync {
+    /// All potential candidates for an argument.
+    ///
+    /// See [`CompletionCandidate`] for more information.
+    fn completions(&self) -> Vec<CompletionCandidate>;
+}
+
+impl<F> CustomCompleter for F
+where
+    F: Fn() -> Vec<CompletionCandidate> + Send + Sync,
+{
+    fn completions(&self) -> Vec<CompletionCandidate> {
+        self()
+    }
+}
+
+/// A wrapper for custom completer
+///
+/// # Example
+///
+/// ```rust
+/// use clap::Parser;
+/// use clap_complete::dynamic::{ArgValueCompleter, CompletionCandidate};
+///
+/// #[derive(Debug, Parser)]
+/// struct Cli {
+///     #[arg(long, add = ArgValueCompleter::new(|| { vec![
+///         CompletionCandidate::new("foo"),
+///         CompletionCandidate::new("bar"),
+///         CompletionCandidate::new("baz")] }))]
+///     custom: Option<String>,
+/// }
+///    
+/// ```
+#[derive(Clone)]
+pub struct ArgValueCompleter(Arc<dyn CustomCompleter>);
+
+impl ArgValueCompleter {
+    /// Create a new `ArgValueCompleter` with a custom completer
+    pub fn new<C: CustomCompleter>(completer: C) -> Self
+    where
+        C: 'static + CustomCompleter,
+    {
+        Self(Arc::new(completer))
+    }
+}
+
+impl std::fmt::Debug for ArgValueCompleter {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_str(type_name::<Self>())
+    }
+}
+
+impl ArgExt for ArgValueCompleter {}
diff --git a/clap_complete/tests/testsuite/dynamic.rs b/clap_complete/tests/testsuite/dynamic.rs
index b909159134b..31214fa1339 100644
--- a/clap_complete/tests/testsuite/dynamic.rs
+++ b/clap_complete/tests/testsuite/dynamic.rs
@@ -4,6 +4,7 @@ use std::fs;
 use std::path::Path;
 
 use clap::{builder::PossibleValue, Command};
+use clap_complete::dynamic::{ArgValueCompleter, CompletionCandidate, CustomCompleter};
 use snapbox::assert_data_eq;
 
 macro_rules! complete {
@@ -592,9 +593,33 @@ val3
 
 #[test]
 fn suggest_custom_arg_value() {
-    let mut cmd = Command::new("dynamic").arg(clap::Arg::new("custom").long("custom"));
+    #[derive(Debug)]
+    struct MyCustomCompleter {}
+
+    impl CustomCompleter for MyCustomCompleter {
+        fn completions(&self) -> Vec<CompletionCandidate> {
+            vec![
+                CompletionCandidate::new("custom1"),
+                CompletionCandidate::new("custom2"),
+                CompletionCandidate::new("custom3"),
+            ]
+        }
+    }
+
+    let mut cmd = Command::new("dynamic").arg(
+        clap::Arg::new("custom")
+            .long("custom")
+            .add::<ArgValueCompleter>(ArgValueCompleter::new(MyCustomCompleter {})),
+    );
 
-    assert_data_eq!(complete!(cmd, "--custom [TAB]"), snapbox::str![""],);
+    assert_data_eq!(
+        complete!(cmd, "--custom [TAB]"),
+        snapbox::str![
+            "custom1
+custom2
+custom3"
+        ],
+    );
 }
 
 #[test]