From 0b7fc4543fafcc74725270ed145ef96a0b0aa8df Mon Sep 17 00:00:00 2001
From: David Arnold <dgx.arnold@gmail.com>
Date: Fri, 6 Aug 2021 11:36:05 -0500
Subject: [PATCH] Add initial completions support

---
 Cargo.lock | 12 ++++++++++++
 Cargo.toml |  1 +
 src/cli.rs | 42 ++++++++++++++++++++++++++++++++++++++----
 3 files changed, 51 insertions(+), 4 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 08ac3bfc..ff09b72f 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1,5 +1,7 @@
 # This file is automatically @generated by Cargo.
 # It is not intended for manual editing.
+version = 3
+
 [[package]]
 name = "aho-corasick"
 version = "0.7.15"
@@ -110,6 +112,15 @@ dependencies = [
  "syn",
 ]
 
+[[package]]
+name = "clap_generate"
+version = "3.0.0-beta.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "adf420f8b687b628d2915ccfd43a660c437a170432e3fbcb66944e8717a0d68f"
+dependencies = [
+ "clap",
+]
+
 [[package]]
 name = "crossbeam-channel"
 version = "0.4.4"
@@ -136,6 +147,7 @@ name = "deploy-rs"
 version = "0.1.0"
 dependencies = [
  "clap",
+ "clap_generate",
  "flexi_logger",
  "fork",
  "futures-util",
diff --git a/Cargo.toml b/Cargo.toml
index 0ded1259..8ae5a230 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -12,6 +12,7 @@ edition = "2018"
 
 [dependencies]
 clap = "3.0.0-beta.2"
+clap_generate = "3.0.0-beta.2"
 flexi_logger = "0.16"
 fork = "0.1"
 futures-util = "0.3.6"
diff --git a/src/cli.rs b/src/cli.rs
index a8592a69..30e79ef5 100644
--- a/src/cli.rs
+++ b/src/cli.rs
@@ -6,7 +6,9 @@
 use std::collections::HashMap;
 use std::io::{stdin, stdout, Write};
 
-use clap::{Clap, ArgMatches, FromArgMatches};
+use clap::{App, ArgEnum, Clap, IntoApp, ValueHint, ArgMatches, FromArgMatches};
+use clap_generate::generators::{Bash, Elvish, Fish, PowerShell, Zsh};
+use clap_generate::{generate, Generator};
 
 use crate as deploy;
 
@@ -18,16 +20,31 @@ use std::process::Stdio;
 use thiserror::Error;
 use tokio::process::Command;
 
+
+#[derive(ArgEnum, Debug, Clone, PartialEq)]
+pub enum GeneratorChoice {
+    Bash,
+    Elvish,
+    Fish,
+    #[clap(name = "powershell")]
+    PowerShell,
+    Zsh,
+}
+
 /// Simple Rust rewrite of a simple Nix Flake deployment tool
 #[derive(Clap, Debug, Clone)]
-#[clap(version = "1.0", author = "Serokell <https://serokell.io/>")]
+#[clap(name = "deploy", version = "1.0", author = "Serokell <https://serokell.io/>")]
 pub struct Opts {
+    /// If provided, outputs the completion file for given shell
+    #[clap(long = "generate", arg_enum)]
+    generator: Option<GeneratorChoice>,
+
     /// The flake to deploy
-    #[clap(group = "deploy")]
+    #[clap(group = "deploy", value_hint = ValueHint::DirPath)]
     target: Option<String>,
 
     /// A list of flakes to deploy alternatively
-    #[clap(long, group = "deploy")]
+    #[clap(long, group = "deploy", value_hint = ValueHint::DirPath)]
     targets: Option<Vec<String>>,
     /// Check signatures when using `nix copy`
     #[clap(short, long)]
@@ -91,6 +108,10 @@ pub struct Opts {
     rollback_succeeded: Option<bool>,
 }
 
+fn print_completions<G: Generator>(app: &mut App) {
+    generate::<G, _>(app, app.get_name().to_string(), &mut stdout());
+}
+
 /// Returns if the available Nix installation supports flakes
 async fn test_flake_support() -> Result<bool, std::io::Error> {
     debug!("Checking for flake support");
@@ -619,6 +640,19 @@ pub async fn run(args: Option<&ArgMatches>) -> Result<(), RunError> {
         deploy::LoggerType::Deploy,
     )?;
 
+    if let Some(generator) = opts.generator {
+        let mut app = Opts::into_app();
+        info!("Generating completion file for {:?}...", generator);
+        match generator {
+            GeneratorChoice::Bash => print_completions::<Bash>(&mut app),
+            GeneratorChoice::Elvish => print_completions::<Elvish>(&mut app),
+            GeneratorChoice::Fish => print_completions::<Fish>(&mut app),
+            GeneratorChoice::PowerShell => print_completions::<PowerShell>(&mut app),
+            GeneratorChoice::Zsh => print_completions::<Zsh>(&mut app),
+        };
+        return Ok(())
+    }
+
     fn maybe_default_target(target: Option<String>) -> Vec<String> {
         match (target, std::env::var("DEPLOY_RS_DEFAULT_FLAKE_ROOT"), std::env::var("DEPLOY_RS_DEFAULT_NODE")) {
             (None, _, _) => vec![".".to_string()],