From 245cba27db626272e2df61812a4e2177f89387b1 Mon Sep 17 00:00:00 2001
From: shannmu <shanmu1901@gmail.com>
Date: Tue, 17 Sep 2024 23:06:50 +0800
Subject: [PATCH] feat: Add custom completer for completing registry name

---
 src/bin/cargo/commands/add.rs     |  6 +++++-
 src/cargo/util/command_prelude.rs | 22 +++++++++++++++++++++-
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/src/bin/cargo/commands/add.rs b/src/bin/cargo/commands/add.rs
index 4c9436df696..263b29b1fac 100644
--- a/src/bin/cargo/commands/add.rs
+++ b/src/bin/cargo/commands/add.rs
@@ -144,7 +144,11 @@ This is the catch all, handling hashes to named references in remote repositorie
                 .long("registry")
                 .action(ArgAction::Set)
                 .value_name("NAME")
-                .help("Package registry for this dependency"),
+                .help("Package registry for this dependency")
+                .add(clap_complete::ArgValueCandidates::new(|| {
+                    let candidates = get_registry_candidates();
+                    candidates.unwrap_or_default()
+                })),
         ])
         .next_help_heading("Section")
         .args([
diff --git a/src/cargo/util/command_prelude.rs b/src/cargo/util/command_prelude.rs
index 247143319f9..16e588cd480 100644
--- a/src/cargo/util/command_prelude.rs
+++ b/src/cargo/util/command_prelude.rs
@@ -386,7 +386,12 @@ pub trait CommandExt: Sized {
     }
 
     fn arg_registry(self, help: &'static str) -> Self {
-        self._arg(opt("registry", help).value_name("REGISTRY"))
+        self._arg(opt("registry", help).value_name("REGISTRY").add(
+            clap_complete::ArgValueCandidates::new(|| {
+                let candidates = get_registry_candidates();
+                candidates.unwrap_or_default()
+            }),
+        ))
     }
 
     fn arg_index(self, help: &'static str) -> Self {
@@ -1063,6 +1068,21 @@ pub fn lockfile_path(
     return Ok(Some(path));
 }
 
+pub fn get_registry_candidates() -> CargoResult<Vec<clap_complete::CompletionCandidate>> {
+    let gctx = new_gctx_for_completions()?;
+
+    if let Ok(Some(registries)) =
+        gctx.get::<Option<HashMap<String, HashMap<String, String>>>>("registries")
+    {
+        Ok(registries
+            .keys()
+            .map(|name| clap_complete::CompletionCandidate::new(name.to_owned()))
+            .collect())
+    } else {
+        Ok(vec![])
+    }
+}
+
 fn get_example_candidates() -> Vec<clap_complete::CompletionCandidate> {
     get_targets_from_metadata()
         .unwrap_or_default()