Merge branch 'add-shell-completion'

1 files changed, 50 insertions, 12 deletions

diff --git a/mullvad-cli/src/main.rs b/mullvad-cli/src/main.rs
index b84307c475..3d8d2a4aef 100644
--- a/mullvad-cli/src/main.rs
+++ b/mullvad-cli/src/main.rs
@@ -1,13 +1,14 @@
 #![deny(rust_2018_idioms)]
 
-use clap::{crate_authors, crate_description, crate_name};
+use clap::{crate_authors, crate_description};
 use mullvad_ipc_client::{new_standalone_ipc_client, DaemonRpcClient};
-use std::io;
+use std::{collections::HashMap, io};
 use talpid_types::ErrorExt;
 
 mod cmds;
 mod location;
 
+pub const BIN_NAME: &str = "mullvad";
 pub const PRODUCT_VERSION: &str = include_str!(concat!(env!("OUT_DIR"), "/product-version.txt"));
 
 pub type Result<T> = std::result::Result<T, Error>;
@@ -50,8 +51,53 @@ fn run() -> Result<()> {
     env_logger::init();
 
     let commands = cmds::get_commands();
+    let app = build_cli(&commands);
 
-    let app = clap::App::new(crate_name!())
+    #[cfg(feature = "shell-completions")]
+    let app = app.subcommand(
+        clap::SubCommand::with_name("shell-completions")
+            .about("Generates completion scripts for your shell")
+            .arg(
+                clap::Arg::with_name("SHELL")
+                    .required(true)
+                    .possible_values(&clap::Shell::variants()[..])
+                    .help("The shell to generate the script for"),
+            )
+            .arg(
+                clap::Arg::with_name("DIR")
+                    .default_value("./")
+                    .help("Output directory where the shell completions are written"),
+            ),
+    );
+
+    let app_matches = app.get_matches();
+    match app_matches.subcommand() {
+        #[cfg(feature = "shell-completions")]
+        ("shell-completions", Some(sub_matches)) => {
+            let shell = sub_matches
+                .value_of("SHELL")
+                .unwrap()
+                .parse()
+                .expect("Invalid shell");
+            let out_dir = sub_matches.value_of_os("DIR").unwrap();
+            build_cli(&commands).gen_completions(BIN_NAME, shell, out_dir);
+            Ok(())
+        }
+        (sub_name, Some(sub_matches)) => {
+            if let Some(cmd) = commands.get(sub_name) {
+                cmd.run(sub_matches)
+            } else {
+                unreachable!("No command matched");
+            }
+        }
+        (_, None) => {
+            unreachable!("No subcommand matches");
+        }
+    }
+}
+
+fn build_cli(commands: &HashMap<&'static str, Box<dyn Command>>) -> clap::App<'static, 'static> {
+    clap::App::new(BIN_NAME)
         .version(PRODUCT_VERSION)
         .author(crate_authors!())
         .about(crate_description!())
@@ -60,15 +106,7 @@ fn run() -> Result<()> {
             clap::AppSettings::DisableHelpSubcommand,
             clap::AppSettings::VersionlessSubcommands,
         ])
-        .subcommands(commands.values().map(|cmd| cmd.clap_subcommand()));
-
-    let app_matches = app.get_matches();
-    let (subcommand_name, subcommand_matches) = app_matches.subcommand();
-    if let Some(cmd) = commands.get(subcommand_name) {
-        cmd.run(subcommand_matches.expect("No command matched"))
-    } else {
-        unreachable!("No command matched");
-    }
+        .subcommands(commands.values().map(|cmd| cmd.clap_subcommand()))
 }
 
 pub trait Command {