From 779c1cc7103cd3ce29d4e2e8e6460e3b30786a0e Mon Sep 17 00:00:00 2001
From: Daylin Morgan <me@dayl.in>
Date: Wed, 27 Mar 2024 11:54:33 -0500
Subject: [PATCH] add support for zsh completions

---
 pkgs/oizys/oizys-rs/Cargo.lock  | 10 ++++++++
 pkgs/oizys/oizys-rs/Cargo.toml  |  1 +
 pkgs/oizys/oizys-rs/default.nix | 12 +++++++++-
 pkgs/oizys/oizys-rs/src/main.rs | 42 +++++++++++++++++++++++++--------
 4 files changed, 54 insertions(+), 11 deletions(-)

diff --git a/pkgs/oizys/oizys-rs/Cargo.lock b/pkgs/oizys/oizys-rs/Cargo.lock
index 69d3886..8e20ee4 100644
--- a/pkgs/oizys/oizys-rs/Cargo.lock
+++ b/pkgs/oizys/oizys-rs/Cargo.lock
@@ -130,6 +130,15 @@ dependencies = [
  "strsim",
 ]
 
+[[package]]
+name = "clap_complete"
+version = "4.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "885e4d7d5af40bfb99ae6f9433e292feac98d452dcb3ec3d25dfe7552b77da8c"
+dependencies = [
+ "clap",
+]
+
 [[package]]
 name = "clap_derive"
 version = "4.5.4"
@@ -372,6 +381,7 @@ name = "oizys"
 version = "0.1.0"
 dependencies = [
  "clap",
+ "clap_complete",
  "homedir",
  "hostname",
 ]
diff --git a/pkgs/oizys/oizys-rs/Cargo.toml b/pkgs/oizys/oizys-rs/Cargo.toml
index 8f77134..6a14554 100644
--- a/pkgs/oizys/oizys-rs/Cargo.toml
+++ b/pkgs/oizys/oizys-rs/Cargo.toml
@@ -7,5 +7,6 @@ edition = "2021"
 
 [dependencies]
 clap = { version = "4.5.4", features = ["derive"] }
+clap_complete = "4.5.1"
 homedir = "0.2.1"
 hostname = "0.3.1"
diff --git a/pkgs/oizys/oizys-rs/default.nix b/pkgs/oizys/oizys-rs/default.nix
index a35873f..4f32b66 100644
--- a/pkgs/oizys/oizys-rs/default.nix
+++ b/pkgs/oizys/oizys-rs/default.nix
@@ -1,4 +1,7 @@
-{rustPlatform}:
+{
+  installShellFiles,
+  rustPlatform,
+}:
 rustPlatform.buildRustPackage {
   pname = "oizys";
   version = "unstable";
@@ -6,4 +9,11 @@ rustPlatform.buildRustPackage {
   cargoLock = {
     lockFile = ./Cargo.lock;
   };
+
+  nativeBuildInputs = [installShellFiles];
+
+  postInstall = ''
+    installShellCompletion --cmd oizys \
+      --zsh <($out/bin/oizys --completions zsh)
+  '';
 }
diff --git a/pkgs/oizys/oizys-rs/src/main.rs b/pkgs/oizys/oizys-rs/src/main.rs
index cf33c11..b3bcd0b 100644
--- a/pkgs/oizys/oizys-rs/src/main.rs
+++ b/pkgs/oizys/oizys-rs/src/main.rs
@@ -1,5 +1,6 @@
-use clap::{Parser, Subcommand};
-use std::{env, path::PathBuf, process::Command};
+use clap::{CommandFactory, Parser, Subcommand};
+use clap_complete::{generate, Generator, Shell};
+use std::{env, io, path::PathBuf, process::Command};
 
 #[derive(Parser)]
 #[command(version, about, long_about = None)]
@@ -20,8 +21,12 @@ struct Cli {
     #[arg(long, global=true, action = clap::ArgAction::SetTrue)]
     no_pinix: bool,
 
+    /// generate shell completion
+    #[arg(long, value_enum)]
+    completions: Option<Shell>,
+
     #[command(subcommand)]
-    command: Commands,
+    command: Option<Commands>,
 }
 
 #[derive(Debug, Subcommand)]
@@ -49,6 +54,10 @@ enum Commands {
     Path {},
 }
 
+fn print_completions<G: Generator>(gen: G, cmd: &mut clap::Command) {
+    generate(gen, cmd, cmd.get_name().to_string(), &mut io::stdout());
+}
+
 #[derive(Debug)]
 struct Oizys {
     host: String,
@@ -148,6 +157,13 @@ fn main() {
     let cli = Cli::parse();
     let oizys = Oizys::new(cli.host, cli.flake, cli.no_pinix, cli.verbose);
 
+    if let Some(completions) = cli.completions {
+        let mut cmd = Cli::command();
+        eprintln!("Generating completion for {completions:?}");
+        print_completions(completions, &mut cmd);
+        std::process::exit(0);
+    }
+
     if oizys.verbose > 2 {
         println!("-vv is max verbosity")
     }
@@ -155,12 +171,18 @@ fn main() {
         println!("{:?}", oizys)
     }
 
-    match &cli.command {
-        Commands::Dry {} => oizys.build(true),
-        Commands::Build {} => oizys.build(false),
-        Commands::Path {} => println!("{}", oizys.output()),
-        Commands::Boot {} => oizys.nixos_rebuild("boot"),
-        Commands::Switch {} => oizys.nixos_rebuild("switch"),
-        Commands::Cache { name } => oizys.cache(name),
+    if let Some(command) = &cli.command {
+        match command {
+            Commands::Dry {} => oizys.build(true),
+            Commands::Build {} => oizys.build(false),
+            Commands::Path {} => println!("{}", oizys.output()),
+            Commands::Boot {} => oizys.nixos_rebuild("boot"),
+            Commands::Switch {} => oizys.nixos_rebuild("switch"),
+            Commands::Cache { name } => oizys.cache(name),
+        }
+    } else {
+        eprintln!("No subcommand provided..");
+        let mut cmd = Cli::command();
+        cmd.print_help().unwrap();
     }
 }