Skip to main content

zql_cli/util/
complete.rs

1use crate::core::system::System;
2use crate::db::config::Config;
3use crate::error::MyResult;
4use crate::regex;
5use clap::Command;
6use clap_complete::Shell;
7use std::io::{Read, Write};
8
9pub fn generate_completion<F: Read + Write, S: System<F>, W: Write>(
10    system: S,
11    writer: &mut W,
12    command: &mut Command,
13    shell: Shell,
14) -> MyResult<()> {
15    if shell == Shell::Bash {
16        let mut buffer = Vec::new();
17        clap_complete::generate(shell, command, String::from("zql"), &mut buffer);
18        #[cfg(windows)]
19        clap_complete::generate(shell, command, String::from("zql.exe"), &mut buffer);
20        let aliases = Config::new(system)?.with_config()?.get_aliases();
21        let buffer = transform_completion(buffer, aliases);
22        writer.write_all(&buffer)?;
23        return Ok(());
24    }
25    clap_complete::generate(shell, command, String::from("zql"), writer);
26    Ok(())
27}
28
29fn transform_completion(buffer: Vec<u8>, aliases: Vec<String>) -> Vec<u8> {
30    let mut lines = String::from_utf8(buffer)
31        .unwrap_or_default()
32        .lines()
33        .map(str::to_string)
34        .collect::<Vec<_>>();
35    let alias_regex = regex!(r"<ALIAS>");
36    let alias_replace = aliases.join(" ");
37    let compgen_regex = regex!(r"\bcompgen -f\b");
38    let compgen_replace = format!("compgen -W \"{}\" --", alias_replace);
39    let mut found = false;
40    for index in 0..lines.len() {
41        lines[index] = alias_regex.replace(&lines[index], &alias_replace).to_string();
42        if lines[index].trim() == "--config)" {
43            found = true;
44        } else if found {
45            lines[index] = compgen_regex.replace(&lines[index], &compgen_replace).to_string();
46            found = false;
47        }
48    }
49    let mut buffer = lines.join("\n");
50    buffer.push_str("\n");
51    buffer.into_bytes()
52}
53
54#[cfg(test)]
55mod tests {
56    use crate::cli::Cli;
57    use crate::core::system::tests::MockSystem;
58    use crate::error::MyResult;
59    use crate::util::complete::generate_completion;
60    use clap::CommandFactory;
61    use clap_complete::Shell;
62    use itertools::Itertools;
63    use pretty_assertions::assert_eq;
64
65    #[test]
66    fn test_bash_completion_is_generated_for_zero_drivers() -> MyResult<()> {
67        let lines = get_lines(Shell::Bash, "")?;
68        let add_command = find_line(&lines, "zql__config__add)");
69        let remove_command = find_line(&lines, "zql__config__remove)");
70        let clear_command = find_line(&lines, "zql__config__clear)");
71        let config_option = find_line(&lines, "--config)");
72        assert_eq!(add_command.contains("--version  [ODBC]\""), true, "{}", add_command);
73        assert_eq!(remove_command.contains("--version \""), true, "{}", remove_command);
74        assert_eq!(clear_command.contains("--version\""), true, "{}", clear_command);
75        assert_eq!(config_option, "COMPREPLY=($(compgen -W \"\" -- \"${cur}\"))");
76        Ok(())
77    }
78
79    #[test]
80    fn test_bash_completion_is_generated_for_three_drivers() -> MyResult<()> {
81        let source = "\
82{
83    \"dsn\": { \"default\": false, \"odbc\": { } },
84    \"mysql\": { \"default\": false, \"odbc\": { } },
85    \"sqlite\": { \"default\": false, \"odbc\": { } }
86}
87";
88        let lines = get_lines(Shell::Bash, source)?;
89        let add_command = find_line(&lines, "zql__config__add)");
90        let remove_command = find_line(&lines, "zql__config__remove)");
91        let clear_command = find_line(&lines, "zql__config__clear)");
92        let config_option = find_line(&lines, "--config)");
93        assert_eq!(add_command.contains("--version dsn mysql sqlite [ODBC]\""), true, "{}", add_command);
94        assert_eq!(remove_command.contains("--version dsn mysql sqlite\""), true, "{}", remove_command);
95        assert_eq!(clear_command.contains("--version\""), true, "{}", clear_command);
96        assert_eq!(config_option, "COMPREPLY=($(compgen -W \"dsn mysql sqlite\" -- \"${cur}\"))");
97        Ok(())
98    }
99
100    fn get_lines(shell: Shell, source: &str) -> MyResult<Vec<String>> {
101        let system = MockSystem::new(source);
102        let mut buffer = Vec::new();
103        let mut command = Cli::command();
104        generate_completion(system, &mut buffer, &mut command, shell)?;
105        let lines = String::from_utf8(buffer)
106            .unwrap_or_default()
107            .lines()
108            .map(str::trim)
109            .map(str::to_string)
110            .collect();
111        Ok(lines)
112    }
113
114    fn find_line(lines: &[String], text: &str) -> String {
115        for (first, second) in lines.iter().tuple_windows() {
116            if first.trim() == text {
117                return second.trim().to_string()
118            }
119        }
120        String::new()
121    }
122}