1 use std::io::Write;
2 
3 use clap::builder::StyledStr;
4 use clap::{Arg, Command};
5 
6 use crate::generator::{utils, Generator};
7 use crate::INTERNAL_ERROR_MSG;
8 
9 /// Generate powershell completion file
10 #[derive(Copy, Clone, PartialEq, Eq, Debug)]
11 pub struct PowerShell;
12 
13 impl Generator for PowerShell {
file_name(&self, name: &str) -> String14     fn file_name(&self, name: &str) -> String {
15         format!("_{name}.ps1")
16     }
17 
generate(&self, cmd: &Command, buf: &mut dyn Write)18     fn generate(&self, cmd: &Command, buf: &mut dyn Write) {
19         let bin_name = cmd
20             .get_bin_name()
21             .expect("crate::generate should have set the bin_name");
22 
23         let subcommands_cases = generate_inner(cmd, "");
24 
25         let result = format!(
26             r#"
27 using namespace System.Management.Automation
28 using namespace System.Management.Automation.Language
29 
30 Register-ArgumentCompleter -Native -CommandName '{bin_name}' -ScriptBlock {{
31     param($wordToComplete, $commandAst, $cursorPosition)
32 
33     $commandElements = $commandAst.CommandElements
34     $command = @(
35         '{bin_name}'
36         for ($i = 1; $i -lt $commandElements.Count; $i++) {{
37             $element = $commandElements[$i]
38             if ($element -isnot [StringConstantExpressionAst] -or
39                 $element.StringConstantType -ne [StringConstantType]::BareWord -or
40                 $element.Value.StartsWith('-') -or
41                 $element.Value -eq $wordToComplete) {{
42                 break
43         }}
44         $element.Value
45     }}) -join ';'
46 
47     $completions = @(switch ($command) {{{subcommands_cases}
48     }})
49 
50     $completions.Where{{ $_.CompletionText -like "$wordToComplete*" }} |
51         Sort-Object -Property ListItemText
52 }}
53 "#
54         );
55 
56         w!(buf, result.as_bytes());
57     }
58 }
59 
60 // Escape string inside single quotes
escape_string(string: &str) -> String61 fn escape_string(string: &str) -> String {
62     string.replace('\'', "''")
63 }
64 
escape_help<T: ToString>(help: Option<&StyledStr>, data: T) -> String65 fn escape_help<T: ToString>(help: Option<&StyledStr>, data: T) -> String {
66     match help {
67         Some(help) => escape_string(&help.to_string().replace('\n', " ")),
68         _ => data.to_string(),
69     }
70 }
71 
generate_inner(p: &Command, previous_command_name: &str) -> String72 fn generate_inner(p: &Command, previous_command_name: &str) -> String {
73     debug!("generate_inner");
74 
75     let command_name = if previous_command_name.is_empty() {
76         p.get_bin_name().expect(INTERNAL_ERROR_MSG).to_string()
77     } else {
78         format!("{};{}", previous_command_name, &p.get_name())
79     };
80 
81     let mut completions = String::new();
82     let preamble = String::from("\n            [CompletionResult]::new(");
83 
84     for option in p.get_opts() {
85         generate_aliases(&mut completions, &preamble, option);
86     }
87 
88     for flag in utils::flags(p) {
89         generate_aliases(&mut completions, &preamble, &flag);
90     }
91 
92     for subcommand in p.get_subcommands() {
93         let data = &subcommand.get_name();
94         let tooltip = escape_help(subcommand.get_about(), data);
95 
96         completions.push_str(&preamble);
97         completions.push_str(
98             format!("'{data}', '{data}', [CompletionResultType]::ParameterValue, '{tooltip}')")
99                 .as_str(),
100         );
101     }
102 
103     let mut subcommands_cases = format!(
104         r"
105         '{}' {{{}
106             break
107         }}",
108         &command_name, completions
109     );
110 
111     for subcommand in p.get_subcommands() {
112         let subcommand_subcommands_cases = generate_inner(subcommand, &command_name);
113         subcommands_cases.push_str(&subcommand_subcommands_cases);
114     }
115 
116     subcommands_cases
117 }
118 
generate_aliases(completions: &mut String, preamble: &String, arg: &Arg)119 fn generate_aliases(completions: &mut String, preamble: &String, arg: &Arg) {
120     use std::fmt::Write as _;
121 
122     if let Some(aliases) = arg.get_short_and_visible_aliases() {
123         let tooltip = escape_help(arg.get_help(), aliases[0]);
124         for alias in aliases {
125             let _ = write!(
126                 completions,
127                 "{preamble}'-{alias}', '{alias}{}', [CompletionResultType]::ParameterName, '{tooltip}')",
128                 // make PowerShell realize there is a difference between `-s` and `-S`
129                 if alias.is_uppercase() { " " } else { "" },
130             );
131         }
132     }
133     if let Some(aliases) = arg.get_long_and_visible_aliases() {
134         let tooltip = escape_help(arg.get_help(), aliases[0]);
135         for alias in aliases {
136             let _ = write!(
137                 completions,
138                 "{preamble}'--{alias}', '{alias}', [CompletionResultType]::ParameterName, '{tooltip}')"
139             );
140         }
141     }
142 }
143