From e9b8de0018b3dd91d72db6e3c164aa18a1fe17d9 Mon Sep 17 00:00:00 2001
From: Karim Abou Zeid <7303830+kabouzeid@users.noreply.github.com>
Date: Mon, 17 Apr 2023 23:33:25 +0200
Subject: [PATCH] feat: add shell completions

---
 Cargo.lock  | 10 ++++++++++
 Cargo.toml  |  1 +
 README.md   | 33 +++++++++++++++++++++++++--------
 src/main.rs | 25 ++++++++++++++++++++++++-
 4 files changed, 60 insertions(+), 9 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 9c38e8b..ef30994 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -114,6 +114,15 @@ dependencies = [
  "strsim",
 ]
 
+[[package]]
+name = "clap_complete"
+version = "4.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "01c22dcfb410883764b29953103d9ef7bb8fe21b3fa1158bc99986c2067294bd"
+dependencies = [
+ "clap",
+]
+
 [[package]]
 name = "clap_derive"
 version = "4.2.0"
@@ -602,6 +611,7 @@ name = "turm"
 version = "0.2.0"
 dependencies = [
  "clap",
+ "clap_complete",
  "crossbeam",
  "crossterm",
  "lazy_static",
diff --git a/Cargo.toml b/Cargo.toml
index 2bb45f4..8fc3097 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -11,6 +11,7 @@ edition = "2021"
 
 [dependencies]
 clap = { version = "4.2.2", features = ["derive"] }
+clap_complete = "4.2.0"
 crossbeam = "0.8.2"
 crossterm = "0.25.0"
 lazy_static = "1.4.0"
diff --git a/README.md b/README.md
index c2f7344..80dc539 100644
--- a/README.md
+++ b/README.md
@@ -4,10 +4,13 @@ A text-based user interface (TUI) for the [Slurm Workload Manager](https://slurm
 
 <img alt="turm demo" src="https://user-images.githubusercontent.com/7303830/228503846-3e5abc04-2c1e-422e-844b-d12ca097403a.gif" width="100%" />
 
-`turm` accepts the same CLI flags as `squeue` (see [man squeue](https://slurm.schedmd.com/squeue.html#SECTION_OPTIONS)), and passes them on when querying the jobs.
-```bash
-turm [--user <username>] [--partition <partition>] [...]
-```
+`turm` accepts the same options as `squeue` (see [man squeue](https://slurm.schedmd.com/squeue.html#SECTION_OPTIONS)). Use `turm --help` to get a list of all available options.
+
+## Project status
+
+This project is currently a work in progress but the basic functionality is there.
+There are still some missing features that need to be implemented and a few visual bugs present at the moment. 
+Please feel free to submit any issues or feedback you may have.
 
 ## Installation
 
@@ -24,11 +27,25 @@ cargo install --path .
 The [release page](https://github.com/kabouzeid/turm/releases) includes precompiled binaries for Linux, macOS and Windows.
 Statically-linked binaries are also available: look for archives with `musl` in the file name.
 
-## Project status
+## Shell Completion
 
-This project is currently a work in progress but the basic functionality is there.
-There are still some missing features that need to be implemented and a few visual bugs present at the moment. 
-Please feel free to submit any issues or feedback you may have.
+### Bash
+
+In your `.bashrc`, add the following line:
+```bash
+eval "$(turm completion bash)"
+```
+
+### Zsh
+
+In your `.zshrc`, add the following line:
+```zsh
+eval "$(turm completion zsh)"
+```
+
+### Other Shells
+
+Completion scripts for other shells (`fish`, `elvish` and `powershell`) can be generated with `turm completion <shell>`.
 
 ## How it works
 
diff --git a/src/main.rs b/src/main.rs
index 341fddb..99286b0 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -4,7 +4,10 @@ mod job_watcher;
 mod squeue_args;
 
 use app::App;
+use clap::CommandFactory;
 use clap::Parser;
+use clap::Subcommand;
+use clap_complete::{generate, Shell};
 use crossbeam::channel::{unbounded, Sender};
 use crossterm::{
     event::{self, DisableMouseCapture, EnableMouseCapture, Event},
@@ -18,7 +21,7 @@ use tui::{
     Terminal,
 };
 
-#[derive(Parser, Debug)]
+#[derive(Parser)]
 #[command(author, version, about, long_about = None)]
 struct Cli {
     /// Refresh rate for the job watcher.
@@ -32,10 +35,30 @@ struct Cli {
     /// squeue arguments
     #[command(flatten)]
     squeue_args: SqueueArgs,
+
+    #[command(subcommand)]
+    command: Option<CliCommand>,
+}
+
+#[derive(Subcommand)]
+enum CliCommand {
+    /// Print shell completion script to stdout.
+    Completion {
+        /// The shell to generate completion for.
+        shell: Shell,
+    },
 }
 
 fn main() -> Result<(), io::Error> {
     let args = Cli::parse();
+    match args.command {
+        Some(CliCommand::Completion { shell }) => {
+            let cmd = &mut Cli::command();
+            generate(shell, cmd, cmd.get_name().to_string(), &mut io::stdout());
+            return Ok(());
+        }
+        None => {}
+    }
 
     // setup terminal
     enable_raw_mode()?;