1 use std::ffi::OsStr;
2 use std::ffi::OsString;
3 
4 use clap::builder::StyledStr;
5 use clap_lex::OsStrExt as _;
6 
7 /// Shell-specific completions
8 pub trait Completer {
9     /// The recommended file name for the registration code
file_name(&self, name: &str) -> String10     fn file_name(&self, name: &str) -> String;
11     /// Register for completions
write_registration( &self, name: &str, bin: &str, completer: &str, buf: &mut dyn std::io::Write, ) -> Result<(), std::io::Error>12     fn write_registration(
13         &self,
14         name: &str,
15         bin: &str,
16         completer: &str,
17         buf: &mut dyn std::io::Write,
18     ) -> Result<(), std::io::Error>;
19     /// Complete the given command
write_complete( &self, cmd: &mut clap::Command, args: Vec<OsString>, current_dir: Option<&std::path::Path>, buf: &mut dyn std::io::Write, ) -> Result<(), std::io::Error>20     fn write_complete(
21         &self,
22         cmd: &mut clap::Command,
23         args: Vec<OsString>,
24         current_dir: Option<&std::path::Path>,
25         buf: &mut dyn std::io::Write,
26     ) -> Result<(), std::io::Error>;
27 }
28 
29 /// Complete the given command
complete( cmd: &mut clap::Command, args: Vec<OsString>, arg_index: usize, current_dir: Option<&std::path::Path>, ) -> Result<Vec<(OsString, Option<StyledStr>)>, std::io::Error>30 pub fn complete(
31     cmd: &mut clap::Command,
32     args: Vec<OsString>,
33     arg_index: usize,
34     current_dir: Option<&std::path::Path>,
35 ) -> Result<Vec<(OsString, Option<StyledStr>)>, std::io::Error> {
36     cmd.build();
37 
38     let raw_args = clap_lex::RawArgs::new(args);
39     let mut cursor = raw_args.cursor();
40     let mut target_cursor = raw_args.cursor();
41     raw_args.seek(
42         &mut target_cursor,
43         clap_lex::SeekFrom::Start(arg_index as u64),
44     );
45     // As we loop, `cursor` will always be pointing to the next item
46     raw_args.next_os(&mut target_cursor);
47 
48     // TODO: Multicall support
49     if !cmd.is_no_binary_name_set() {
50         raw_args.next_os(&mut cursor);
51     }
52 
53     let mut current_cmd = &*cmd;
54     let mut pos_index = 1;
55     let mut is_escaped = false;
56     while let Some(arg) = raw_args.next(&mut cursor) {
57         if cursor == target_cursor {
58             return complete_arg(&arg, current_cmd, current_dir, pos_index, is_escaped);
59         }
60 
61         debug!("complete::next: Begin parsing '{:?}'", arg.to_value_os(),);
62 
63         if let Ok(value) = arg.to_value() {
64             if let Some(next_cmd) = current_cmd.find_subcommand(value) {
65                 current_cmd = next_cmd;
66                 pos_index = 1;
67                 continue;
68             }
69         }
70 
71         if is_escaped {
72             pos_index += 1;
73         } else if arg.is_escape() {
74             is_escaped = true;
75         } else if let Some(_long) = arg.to_long() {
76         } else if let Some(_short) = arg.to_short() {
77         } else {
78             pos_index += 1;
79         }
80     }
81 
82     Err(std::io::Error::new(
83         std::io::ErrorKind::Other,
84         "no completion generated",
85     ))
86 }
87 
complete_arg( arg: &clap_lex::ParsedArg<'_>, cmd: &clap::Command, current_dir: Option<&std::path::Path>, pos_index: usize, is_escaped: bool, ) -> Result<Vec<(OsString, Option<StyledStr>)>, std::io::Error>88 fn complete_arg(
89     arg: &clap_lex::ParsedArg<'_>,
90     cmd: &clap::Command,
91     current_dir: Option<&std::path::Path>,
92     pos_index: usize,
93     is_escaped: bool,
94 ) -> Result<Vec<(OsString, Option<StyledStr>)>, std::io::Error> {
95     debug!(
96         "complete_arg: arg={:?}, cmd={:?}, current_dir={:?}, pos_index={}, is_escaped={}",
97         arg,
98         cmd.get_name(),
99         current_dir,
100         pos_index,
101         is_escaped
102     );
103     let mut completions = Vec::new();
104 
105     if !is_escaped {
106         if let Some((flag, value)) = arg.to_long() {
107             if let Ok(flag) = flag {
108                 if let Some(value) = value {
109                     if let Some(arg) = cmd.get_arguments().find(|a| a.get_long() == Some(flag)) {
110                         completions.extend(
111                             complete_arg_value(value.to_str().ok_or(value), arg, current_dir)
112                                 .into_iter()
113                                 .map(|(os, help)| {
114                                     // HACK: Need better `OsStr` manipulation
115                                     (format!("--{}={}", flag, os.to_string_lossy()).into(), help)
116                                 }),
117                         );
118                     }
119                 } else {
120                     completions.extend(longs_and_visible_aliases(cmd).into_iter().filter_map(
121                         |(f, help)| f.starts_with(flag).then(|| (format!("--{f}").into(), help)),
122                     ));
123                 }
124             }
125         } else if arg.is_escape() || arg.is_stdio() || arg.is_empty() {
126             // HACK: Assuming knowledge of is_escape / is_stdio
127             completions.extend(
128                 longs_and_visible_aliases(cmd)
129                     .into_iter()
130                     .map(|(f, help)| (format!("--{f}").into(), help)),
131             );
132         }
133 
134         if arg.is_empty() || arg.is_stdio() || arg.is_short() {
135             let dash_or_arg = if arg.is_empty() {
136                 "-".into()
137             } else {
138                 arg.to_value_os().to_string_lossy()
139             };
140             // HACK: Assuming knowledge of is_stdio
141             completions.extend(
142                 shorts_and_visible_aliases(cmd)
143                     .into_iter()
144                     // HACK: Need better `OsStr` manipulation
145                     .map(|(f, help)| (format!("{}{}", dash_or_arg, f).into(), help)),
146             );
147         }
148     }
149 
150     if let Some(positional) = cmd
151         .get_positionals()
152         .find(|p| p.get_index() == Some(pos_index))
153     {
154         completions.extend(complete_arg_value(arg.to_value(), positional, current_dir));
155     }
156 
157     if let Ok(value) = arg.to_value() {
158         completions.extend(complete_subcommand(value, cmd));
159     }
160 
161     Ok(completions)
162 }
163 
complete_arg_value( value: Result<&str, &OsStr>, arg: &clap::Arg, current_dir: Option<&std::path::Path>, ) -> Vec<(OsString, Option<StyledStr>)>164 fn complete_arg_value(
165     value: Result<&str, &OsStr>,
166     arg: &clap::Arg,
167     current_dir: Option<&std::path::Path>,
168 ) -> Vec<(OsString, Option<StyledStr>)> {
169     let mut values = Vec::new();
170     debug!("complete_arg_value: arg={arg:?}, value={value:?}");
171 
172     if let Some(possible_values) = possible_values(arg) {
173         if let Ok(value) = value {
174             values.extend(possible_values.into_iter().filter_map(|p| {
175                 let name = p.get_name();
176                 name.starts_with(value)
177                     .then(|| (name.into(), p.get_help().cloned()))
178             }));
179         }
180     } else {
181         let value_os = match value {
182             Ok(value) => OsStr::new(value),
183             Err(value_os) => value_os,
184         };
185         match arg.get_value_hint() {
186             clap::ValueHint::Other => {
187                 // Should not complete
188             }
189             clap::ValueHint::Unknown | clap::ValueHint::AnyPath => {
190                 values.extend(complete_path(value_os, current_dir, |_| true));
191             }
192             clap::ValueHint::FilePath => {
193                 values.extend(complete_path(value_os, current_dir, |p| p.is_file()));
194             }
195             clap::ValueHint::DirPath => {
196                 values.extend(complete_path(value_os, current_dir, |p| p.is_dir()));
197             }
198             clap::ValueHint::ExecutablePath => {
199                 use is_executable::IsExecutable;
200                 values.extend(complete_path(value_os, current_dir, |p| p.is_executable()));
201             }
202             clap::ValueHint::CommandName
203             | clap::ValueHint::CommandString
204             | clap::ValueHint::CommandWithArguments
205             | clap::ValueHint::Username
206             | clap::ValueHint::Hostname
207             | clap::ValueHint::Url
208             | clap::ValueHint::EmailAddress => {
209                 // No completion implementation
210             }
211             _ => {
212                 // Safe-ish fallback
213                 values.extend(complete_path(value_os, current_dir, |_| true));
214             }
215         }
216         values.sort();
217     }
218 
219     values
220 }
221 
complete_path( value_os: &OsStr, current_dir: Option<&std::path::Path>, is_wanted: impl Fn(&std::path::Path) -> bool, ) -> Vec<(OsString, Option<StyledStr>)>222 fn complete_path(
223     value_os: &OsStr,
224     current_dir: Option<&std::path::Path>,
225     is_wanted: impl Fn(&std::path::Path) -> bool,
226 ) -> Vec<(OsString, Option<StyledStr>)> {
227     let mut completions = Vec::new();
228 
229     let current_dir = match current_dir {
230         Some(current_dir) => current_dir,
231         None => {
232             // Can't complete without a `current_dir`
233             return Vec::new();
234         }
235     };
236     let (existing, prefix) = value_os
237         .split_once("\\")
238         .unwrap_or((OsStr::new(""), value_os));
239     let root = current_dir.join(existing);
240     debug!("complete_path: root={root:?}, prefix={prefix:?}");
241     let prefix = prefix.to_string_lossy();
242 
243     for entry in std::fs::read_dir(&root)
244         .ok()
245         .into_iter()
246         .flatten()
247         .filter_map(Result::ok)
248     {
249         let raw_file_name = entry.file_name();
250         if !raw_file_name.starts_with(&prefix) {
251             continue;
252         }
253 
254         if entry.metadata().map(|m| m.is_dir()).unwrap_or(false) {
255             let path = entry.path();
256             let mut suggestion = pathdiff::diff_paths(&path, current_dir).unwrap_or(path);
257             suggestion.push(""); // Ensure trailing `/`
258             completions.push((suggestion.as_os_str().to_owned(), None));
259         } else {
260             let path = entry.path();
261             if is_wanted(&path) {
262                 let suggestion = pathdiff::diff_paths(&path, current_dir).unwrap_or(path);
263                 completions.push((suggestion.as_os_str().to_owned(), None));
264             }
265         }
266     }
267 
268     completions
269 }
270 
271 fn complete_subcommand(value: &str, cmd: &clap::Command) -> Vec<(OsString, Option<StyledStr>)> {
272     debug!(
273         "complete_subcommand: cmd={:?}, value={:?}",
274         cmd.get_name(),
275         value
276     );
277 
278     let mut scs = subcommands(cmd)
279         .into_iter()
280         .filter(|x| x.0.starts_with(value))
281         .map(|x| (OsString::from(&x.0), x.1))
282         .collect::<Vec<_>>();
283     scs.sort();
284     scs.dedup();
285     scs
286 }
287 
288 /// Gets all the long options, their visible aliases and flags of a [`clap::Command`].
289 /// Includes `help` and `version` depending on the [`clap::Command`] settings.
290 fn longs_and_visible_aliases(p: &clap::Command) -> Vec<(String, Option<StyledStr>)> {
291     debug!("longs: name={}", p.get_name());
292 
293     p.get_arguments()
294         .filter_map(|a| {
295             a.get_long_and_visible_aliases().map(|longs| {
296                 longs
297                     .into_iter()
298                     .map(|s| (s.to_string(), a.get_help().cloned()))
299             })
300         })
301         .flatten()
302         .collect()
303 }
304 
305 /// Gets all the short options, their visible aliases and flags of a [`clap::Command`].
306 /// Includes `h` and `V` depending on the [`clap::Command`] settings.
307 fn shorts_and_visible_aliases(p: &clap::Command) -> Vec<(char, Option<StyledStr>)> {
308     debug!("shorts: name={}", p.get_name());
309 
310     p.get_arguments()
311         .filter_map(|a| {
312             a.get_short_and_visible_aliases()
313                 .map(|shorts| shorts.into_iter().map(|s| (s, a.get_help().cloned())))
314         })
315         .flatten()
316         .collect()
317 }
318 
319 /// Get the possible values for completion
320 fn possible_values(a: &clap::Arg) -> Option<Vec<clap::builder::PossibleValue>> {
321     if !a.get_num_args().expect("built").takes_values() {
322         None
323     } else {
324         a.get_value_parser()
325             .possible_values()
326             .map(|pvs| pvs.collect())
327     }
328 }
329 
330 /// Gets subcommands of [`clap::Command`] in the form of `("name", "bin_name")`.
331 ///
332 /// Subcommand `rustup toolchain install` would be converted to
333 /// `("install", "rustup toolchain install")`.
334 fn subcommands(p: &clap::Command) -> Vec<(String, Option<StyledStr>)> {
335     debug!("subcommands: name={}", p.get_name());
336     debug!("subcommands: Has subcommands...{:?}", p.has_subcommands());
337 
338     p.get_subcommands()
339         .map(|sc| (sc.get_name().to_string(), sc.get_about().cloned()))
340         .collect()
341 }
342