Add Zsh Completion Script Generation and Unit Tests

src/lib/cli_parser.rs

@@ -7,6 +7,8 @@
 #[path = "cli_parser_test.rs"]
 mod cli_parser_test;
 
+use crate::completion::generate_completions;
+
 use crate::cli::{
     AUTHOR, DEFAULT_LOG_LEVEL, DEFAULT_OUTPUT_FORMAT, DEFAULT_TASK_NAME, DESCRIPTION, VERSION,
 };
@@ -42,6 +44,11 @@ fn get_args(
         None => None,
     };
 
+    cli_args.completion = match cli_parsed.get_first_value("completion") {
+        Some(value) => Some(value.to_string()),
+        None => None,
+    };
+
     cli_args.cwd = cli_parsed.get_first_value("cwd");
 
     let default_log_level = match global_config.log_level {
@@ -214,6 +221,17 @@ fn add_arguments(spec: CliSpec, default_task_name: &str, default_log_level: &str
                 "PROFILE".to_string(),
             )),
         })
+        .add_argument(Argument {
+            name: "completion".to_string(),
+            key: vec!["--completion".to_string()],
+            argument_occurrence: ArgumentOccurrence::Single,
+            value_type: ArgumentValueType::Single,
+            default_value: None,
+            help: Some(ArgumentHelp::TextAndParam(
+                "Will enable completion for the defined tasks for a given shell".to_string(),
+                "COMPLETION".to_string(),
+            )),
+        })
         .add_argument(Argument {
             name: "cwd".to_string(),
             key: vec!["--cwd".to_string()],
@@ -482,6 +500,10 @@ pub fn parse_args(
         let version_text = cliparser::version(&spec);
         println!("{}", version_text);
         Err(CargoMakeError::ExitCode(std::process::ExitCode::SUCCESS))
+    } else if let Some(shell) = cli_parsed.get_first_value("completion") {
+        // Call the function to generate completions
+        generate_completions(&shell);
+        Err(CargoMakeError::ExitCode(std::process::ExitCode::SUCCESS))
     } else {
         Ok(get_args(
             &cli_parsed,
@@ -512,3 +534,4 @@ fn to_owned_vec(vec_option: Option<&Vec<String>>) -> Option<Vec<String>> {
         None => None,
     }
 }
+

src/lib/cli_test.rs

@@ -18,6 +18,7 @@ fn run_makefile_not_found() {
             profile: None,
             log_level: "error".to_string(),
             disable_color: true,
+            completion: None,
             cwd: None,
             env: None,
             env_file: None,
@@ -57,6 +58,7 @@ fn run_empty_task() {
             profile: None,
             log_level: "error".to_string(),
             disable_color: true,
+            completion: None,
             cwd: None,
             env: None,
             env_file: None,
@@ -96,6 +98,7 @@ fn print_empty_task() {
             profile: None,
             log_level: "error".to_string(),
             disable_color: true,
+            completion: None,
             cwd: None,
             env: None,
             env_file: None,
@@ -135,6 +138,7 @@ fn list_empty_task() {
             profile: None,
             log_level: "error".to_string(),
             disable_color: true,
+            completion: None,
             cwd: None,
             env: None,
             env_file: None,
@@ -174,6 +178,7 @@ fn run_file_and_task() {
             profile: None,
             log_level: "error".to_string(),
             disable_color: true,
+            completion: None,
             cwd: None,
             env: None,
             env_file: None,
@@ -216,6 +221,7 @@ fn run_cwd_with_file() {
             profile: None,
             log_level: "error".to_string(),
             disable_color: true,
+            completion: None,
             cwd: Some("..".to_string()),
             env: None,
             env_file: None,
@@ -256,6 +262,7 @@ fn run_file_not_go_to_project_root() {
             profile: None,
             log_level: "error".to_string(),
             disable_color: true,
+            completion: None,
             cwd: None,
             env: None,
             env_file: None,
@@ -296,6 +303,7 @@ fn run_cwd_go_to_project_root_current_dir() {
             profile: None,
             log_level: "error".to_string(),
             disable_color: true,
+            completion: None,
             cwd: None,
             env: None,
             env_file: None,
@@ -339,6 +347,7 @@ fn run_cwd_go_to_project_root_child_dir() {
             profile: None,
             log_level: "error".to_string(),
             disable_color: true,
+            completion: None,
             cwd: None,
             env: None,
             env_file: None,
@@ -382,6 +391,7 @@ fn run_cwd_task_not_found() {
             profile: None,
             log_level: "error".to_string(),
             disable_color: true,
+            completion: None,
             cwd: Some("..".to_string()),
             env: None,
             env_file: None,

src/lib/completion.rs

@@ -0,0 +1,103 @@
+use std::io::BufRead;
+use std::path::Path;
+use std::{fs, io};
+
+/// # Completions Module
+/// 
+/// This module handles the generation of shell completion scripts for the `cargo-make` tool.
+/// 
+/// ## Functionality
+/// - `generate_completion_zsh`: Generates a Zsh completion script, creates the necessary directory, 
+///   and prompts for overwriting existing files.
+/// 
+/// ## Improvements to Consider
+/// 1. **Modularity**: Separate the completion logic into different modules for different shells
+///    (e.g., Zsh, Bash, Fish) to improve code organization.
+/// 2. **Cross-Platform Support**: Abstract the completion generation into a trait or interface 
+///    to facilitate adding support for other shell types.
+/// 3. **Enhanced Error Handling**: Provide more informative error messages for file operations.
+/// 4. **User Input Handling**: Ensure user input is trimmed and handled correctly.
+/// 5. **Testing**: Implement unit tests to verify the correct behavior of completion generation functions.
+
+#[cfg(test)]
+#[path = "completion_test.rs"]
+mod completion_test;
+
+pub fn generate_completions(shell: &str) {
+    match shell {
+        "zsh" => {
+            if let Err(e) = generate_completion_zsh(None) {
+                eprintln!("Error generating Zsh completions: {}", e);
+            }
+        }
+        _ => {
+            eprintln!("Unsupported shell for completion: {}", shell);
+        }
+    }
+}
+
+// Modify the function to accept an optional input stream
+fn generate_completion_zsh(input: Option<&mut dyn io::Read>) -> Result<(), Box<dyn std::error::Error>> {
+    let home_dir = std::env::var("HOME")?;
+    let zfunc_dir = format!("{}/.zfunc", home_dir);
+    let completion_file = format!("{}/_cargo-make", zfunc_dir);
+
+    if !Path::new(&zfunc_dir).exists() {
+        if let Err(e) = fs::create_dir_all(&zfunc_dir) {
+            eprintln!("Failed to create directory {}: {}", zfunc_dir, e);
+            return Err(Box::new(e));
+        }
+        println!("Created directory: {}", zfunc_dir);
+    }
+
+    if Path::new(&completion_file).exists() {
+        let mut input_str = String::new();
+        let reader: Box<dyn io::Read> = match input {
+            Some(input) => Box::new(input),
+            None => Box::new(io::stdin()),
+        };
+
+        // Create a BufReader to read from the provided input or stdin
+        let mut buf_reader = io::BufReader::new(reader);
+        println!("File {} already exists. Overwrite? (y/n): ", completion_file);
+        buf_reader.read_line(&mut input_str)?;
+
+        if input_str.trim().to_lowercase() != "y" {
+            println!("Aborted overwriting the file.");
+            return Ok(());
+        }
+    }
+
+    let completion_script = r#"
+#compdef cargo make cargo-make
+
+_cargo_make() {
+    local tasks
+    local makefile="Makefile.toml"
+
+    if [[ ! -f $makefile ]]; then
+        return 1
+    fi
+
+    tasks=($(awk -F'[\\[\\.\\]]' '/^\[tasks/ {print $3}' "$makefile"))
+
+    if [[ ${#tasks[@]} -eq 0 ]]; then
+        return 1
+    fi
+
+    _describe -t tasks 'cargo-make tasks' tasks
+}
+
+_cargo_make "$@"
+"#;
+
+    fs::write(&completion_file, completion_script)?;
+    println!("\nWrote tasks completion script to: {}", completion_file);
+
+    println!("To enable Zsh completion, add the following lines to your ~/.zshrc:\n");
+    println!("    fpath=(~/.zfunc $fpath)");
+    println!("    autoload -Uz compinit && compinit");
+    println!("\nThen, restart your terminal or run 'source ~/.zshrc'.");
+
+    Ok(())
+}

src/lib/completion_test.rs

@@ -0,0 +1,90 @@
+use std::fs;
+use std::path::Path;
+use std::io::Cursor;
+
+#[cfg(test)]
+mod tests {
+    use crate::completion::generate_completion_zsh;
+
+    use super::*;
+
+    // Function to clean up test environment by removing the completion file
+    fn cleanup() {
+        let home_dir = std::env::var("HOME").expect("Failed to get HOME");
+        let completion_file = format!("{}/.zfunc/_cargo-make", home_dir);
+        if Path::new(&completion_file).exists() {
+            fs::remove_file(&completion_file).expect("Failed to clean up test file");
+        }
+    }
+
+    #[test]
+    fn test_generate_completion_zsh_overwrite_prompt_yes() {
+        cleanup(); // Clean up before the test
+        let input = b"y\n"; // Simulate user input of 'y'
+        let mut reader = Cursor::new(input);
+
+        let result = generate_completion_zsh(Some(&mut reader));
+        assert!(result.is_ok());
+    }
+
+    #[test]
+    fn test_generate_completion_zsh_overwrite_prompt_no() {
+        cleanup(); // Clean up before the test
+        let input = b"n\n"; // Simulate user input of 'n'
+        let mut reader = Cursor::new(input);
+
+        let result = generate_completion_zsh(Some(&mut reader));
+        assert!(result.is_ok());
+    }
+
+    #[test]
+    fn test_generate_completion_zsh_creates_directory() {
+        cleanup(); // Clean up before the test
+
+        let input = b"y\n"; // Simulate user input of 'y'
+        let mut reader = Cursor::new(input);
+
+        let result = generate_completion_zsh(Some(&mut reader));
+        assert!(result.is_ok(), "Should succeed in generating completions");
+
+        // Check if the directory was created
+        let home_dir = std::env::var("HOME").expect("Failed to get HOME");
+        let zfunc_dir = format!("{}/.zfunc", home_dir);
+        assert!(Path::new(&zfunc_dir).exists(), "The zfunc directory should exist");
+    }
+
+    #[test]
+    fn test_generate_completion_zsh_creates_file() {
+        cleanup(); // Clean up before the test
+
+        let input = b"y\n"; // Simulate user input of 'y'
+        let mut reader = Cursor::new(input);
+
+        let result = generate_completion_zsh(Some(&mut reader));
+        assert!(result.is_ok(), "Should succeed in generating completions");
+
+        // Check if the completion file was created
+        let home_dir = std::env::var("HOME").expect("Failed to get HOME");
+        let completion_file = format!("{}/.zfunc/_cargo-make", home_dir);
+        assert!(Path::new(&completion_file).exists(), "The completion file should exist");
+    }
+
+    #[test]
+    fn test_generate_completion_zsh_overwrite_prompt() {
+        cleanup(); // Clean up before the test
+
+        // Create the directory and file first
+        let input = b"y\n"; // Simulate user input of 'y'
+        let mut reader = Cursor::new(input);
+
+        generate_completion_zsh(Some(&mut reader)).expect("Should succeed in generating completions");
+
+        // Simulate user input for overwrite.
+        let input = b"y\n"; // Simulate user input of 'y' again
+        let mut reader = Cursor::new(input);
+
+        let result = generate_completion_zsh(Some(&mut reader));
+        assert!(result.is_ok(), "Should handle overwrite prompt gracefully");
+    }
+
+}

src/lib/mod.rs

@@ -55,6 +55,7 @@ mod cache;
 pub mod cli;
 pub mod cli_commands;
 pub mod cli_parser;
+pub mod completion;
 mod command;
 mod condition;
 pub mod config;

src/lib/types.rs

@@ -91,6 +91,8 @@ pub struct CliArgs {
     pub log_level: String,
     /// Disables colorful output
     pub disable_color: bool,
+    /// Task completion for given shell
+    pub completion: Option<String>,
     /// Current working directory
     pub cwd: Option<String>,
     /// Environment variables
@@ -141,6 +143,7 @@ impl CliArgs {
             profile: None,
             log_level: "info".to_string(),
             disable_color: false,
+            completion: None,
             cwd: None,
             env: None,
             env_file: None,