From 1f0feb5254faa8e938ba83aff96f2b984bf65477 Mon Sep 17 00:00:00 2001
From: Amaan Qureshi <amaanq12@gmail.com>
Date: Tue, 24 Sep 2024 12:37:31 -0400
Subject: [PATCH] feat: add shell completions

---
 Cargo.lock      | 10 ++++++++++
 Cargo.toml      |  1 +
 cli/Cargo.toml  |  1 +
 cli/src/main.rs | 24 ++++++++++++++++++++++--
 4 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 7dacad72..8b131228 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -231,6 +231,15 @@ dependencies = [
  "strsim",
 ]
 
+[[package]]
+name = "clap_complete"
+version = "4.5.29"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8937760c3f4c60871870b8c3ee5f9b30771f792a7045c48bcbba999d7d6b3b8e"
+dependencies = [
+ "clap",
+]
+
 [[package]]
 name = "clap_derive"
 version = "4.5.18"
@@ -1468,6 +1477,7 @@ dependencies = [
  "anyhow",
  "bstr",
  "clap",
+ "clap_complete",
  "ctor",
  "ctrlc",
  "dirs",
diff --git a/Cargo.toml b/Cargo.toml
index 997f5e6b..544e6149 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -50,6 +50,7 @@ clap = { version = "4.5.18", features = [
   "help",
   "unstable-styles",
 ] }
+clap_complete = "4.5.29"
 ctor = "0.2.8"
 ctrlc = { version = "3.4.5", features = ["termination"] }
 dirs = "5.0.1"
diff --git a/cli/Cargo.toml b/cli/Cargo.toml
index 34fa31db..36c2f9c8 100644
--- a/cli/Cargo.toml
+++ b/cli/Cargo.toml
@@ -29,6 +29,7 @@ anstyle.workspace = true
 anyhow.workspace = true
 bstr.workspace = true
 clap.workspace = true
+clap_complete.workspace = true
 ctor.workspace = true
 ctrlc.workspace = true
 dirs.workspace = true
diff --git a/cli/src/main.rs b/cli/src/main.rs
index 7dfc3dc5..8c47090c 100644
--- a/cli/src/main.rs
+++ b/cli/src/main.rs
@@ -7,6 +7,7 @@ use std::{
 use anstyle::{AnsiColor, Color, Style};
 use anyhow::{anyhow, Context, Result};
 use clap::{crate_authors, Args, Command, FromArgMatches as _, Subcommand};
+use clap_complete::{generate, Shell};
 use glob::glob;
 use regex::Regex;
 use tree_sitter::{ffi, Parser, Point};
@@ -45,6 +46,7 @@ enum Commands {
     Tags(Tags),
     Playground(Playground),
     DumpLanguages(DumpLanguages),
+    Complete(Complete),
 }
 
 #[derive(Args)]
@@ -388,6 +390,18 @@ struct DumpLanguages {
     pub config_path: Option<PathBuf>,
 }
 
+#[derive(Args)]
+#[command(about = "Generate shell completions", alias = "comp")]
+struct Complete {
+    #[arg(
+        long,
+        short,
+        value_enum,
+        help = "The shell to generate completions for"
+    )]
+    pub shell: Shell,
+}
+
 fn main() {
     let result = run();
     if let Err(err) = &result {
@@ -426,9 +440,9 @@ fn run() -> Result<()> {
         .arg_required_else_help(true)
         .disable_help_subcommand(true)
         .disable_colored_help(false);
-    let cli = Commands::augment_subcommands(cli);
+    let mut cli = Commands::augment_subcommands(cli);
 
-    let command = Commands::from_arg_matches(&cli.get_matches())?;
+    let command = Commands::from_arg_matches(&cli.clone().get_matches())?;
 
     let current_dir = env::current_dir().unwrap();
     let mut loader = loader::Loader::new()?;
@@ -1000,6 +1014,12 @@ fn run() -> Result<()> {
                 );
             }
         }
+
+        Commands::Complete(complete_options) => {
+            let name = cli.get_name().to_string();
+            let stdout = &mut std::io::stdout();
+            generate(complete_options.shell, &mut cli, name, stdout);
+        }
     }
 
     Ok(())