From 0927fe9dc1ae913aeaa50d5c8a763a4aff32d808 Mon Sep 17 00:00:00 2001
From: Zhaofeng Li <hello@zhaofeng.li>
Date: Tue, 23 Mar 2021 14:14:04 -0700
Subject: [PATCH] cli: Add hidden command to generate shell autocompletion

---
 default.nix | 14 ++++++++++++++
 src/cli.rs  | 28 +++++++++++++++++++++++++---
 2 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/default.nix b/default.nix
index ba07b03..8f50008 100644
--- a/default.nix
+++ b/default.nix
@@ -16,6 +16,20 @@ in rustPlatform.buildRustPackage {
   };
   cargoSha256 = "1yjaqhv9gd86jq56vsrhv6qv3k5qh2pnc4zyxbi2fm2hdrvy0440";
 
+  postBuild = ''
+    mkdir completions
+    for shell in bash fish zsh; do
+      cargo run --frozen -- gen-completions $shell > completions/$shell
+    done
+  '';
+
+  postInstall = ''
+    mkdir -p "$out/share/"{bash-completion/completions,fish/vendor_completions.d,zsh/site-functions}
+    cp completions/bash $out/share/bash-completion/completions/colmena
+    cp completions/fish $out/share/fish/vendor_completions.d/colmena.fish
+    cp completions/zsh $out/share/zsh/site-functions/_colmena
+  '';
+
   # Recursive Nix is not stable yet
   doCheck = false;
 }
diff --git a/src/cli.rs b/src/cli.rs
index a22af52..012a2ae 100644
--- a/src/cli.rs
+++ b/src/cli.rs
@@ -1,6 +1,6 @@
 //! Global CLI Setup.
 
-use clap::{App, AppSettings, Arg};
+use clap::{App, AppSettings, Arg, ArgMatches, SubCommand};
 use crate::command;
 
 macro_rules! register_command {
@@ -24,7 +24,7 @@ macro_rules! handle_command {
     };
 }
 
-pub fn build_cli() -> App<'static, 'static> {
+pub fn build_cli(include_internal: bool) -> App<'static, 'static> {
     let mut app = App::new("Colmena")
         .bin_name("colmena")
         .version("0.1.0")
@@ -57,6 +57,16 @@ For a sample configuration, see <https://github.com/zhaofengli/colmena>.
             .global(true)
             .takes_value(false));
 
+    if include_internal {
+        app = app.subcommand(SubCommand::with_name("gen-completions")
+            .about("Generate shell auto-completion files (Internal)")
+            .setting(AppSettings::Hidden)
+            .arg(Arg::with_name("shell")
+                .index(1)
+                .required(true)
+                .takes_value(true)));
+    }
+
     register_command!(apply, app);
     register_command!(apply_local, app);
     register_command!(build, app);
@@ -68,7 +78,7 @@ For a sample configuration, see <https://github.com/zhaofengli/colmena>.
 }
 
 pub async fn run() {
-    let mut app = build_cli();
+    let mut app = build_cli(true);
     let matches = app.clone().get_matches();
 
     handle_command!(apply, matches);
@@ -78,6 +88,18 @@ pub async fn run() {
     handle_command!("upload-keys", upload_keys, matches);
     handle_command!(exec, matches);
 
+    if let Some(args) = matches.subcommand_matches("gen-completions") {
+        return gen_completions(args);
+    }
+
     app.print_long_help().unwrap();
     println!();
 }
+
+fn gen_completions(args: &ArgMatches<'_>) {
+    let mut app = build_cli(false);
+    let shell: clap::Shell = args.value_of("shell").unwrap()
+        .parse().unwrap();
+
+    app.gen_completions_to("colmena", shell, &mut std::io::stdout());
+}