arrrg/
lib.rs

1#![doc = include_str!("../README.md")]
2
3use std::collections::HashSet;
4use std::fmt;
5use std::str::FromStr;
6
7use getopts::Fail;
8
9//////////////////////////////////////////// subcommands ///////////////////////////////////////////
10
11/// Error returned when subcommand dispatch cannot select a command.
12#[derive(Clone, Debug, Eq, PartialEq)]
13pub struct SubcommandError {
14    command: Option<String>,
15    expected: Vec<String>,
16}
17
18impl SubcommandError {
19    /// Create a missing-subcommand error.
20    pub fn missing(expected: Vec<String>) -> Self {
21        Self {
22            command: None,
23            expected,
24        }
25    }
26
27    /// Create an unknown-subcommand error.
28    pub fn unknown(command: String, expected: Vec<String>) -> Self {
29        Self {
30            command: Some(command),
31            expected,
32        }
33    }
34
35    /// Return the unknown subcommand, or `None` when no subcommand was provided.
36    pub fn command(&self) -> Option<&str> {
37        self.command.as_deref()
38    }
39
40    /// Return the expected command names.
41    pub fn expected(&self) -> &[String] {
42        &self.expected
43    }
44}
45
46impl fmt::Display for SubcommandError {
47    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
48        if let Some(command) = &self.command {
49            write!(f, "unknown subcommand {command:?}")?;
50        } else {
51            write!(f, "missing subcommand")?;
52        }
53        write_expected_subcommands(f, &self.expected)
54    }
55}
56
57impl std::error::Error for SubcommandError {}
58
59fn write_expected_subcommands(f: &mut fmt::Formatter<'_>, expected: &[String]) -> fmt::Result {
60    if expected.is_empty() {
61        return Ok(());
62    }
63    write!(f, "; expected one of: {}", expected.join(", "))
64}
65
66/// Split a command line's free arguments into `(subcommand, remaining_args)`.
67///
68/// arrrg parsers use `getopts::ParsingStyle::StopAtFirstFree`, so recursive subcommand parsing is
69/// naturally represented as:
70///
71/// 1. parse the current command's options,
72/// 2. split the first free argument as the subcommand name,
73/// 3. parse the remaining free arguments with the selected subcommand's [CommandLine].
74pub fn split_subcommand(
75    free: Vec<String>,
76    command_names: &[&str],
77) -> Result<(String, Vec<String>), SubcommandError> {
78    let mut free = free.into_iter();
79    let expected = command_names
80        .iter()
81        .map(|command| command.to_string())
82        .collect::<Vec<_>>();
83    let Some(command) = free.next() else {
84        return Err(SubcommandError::missing(expected));
85    };
86    if !command_names.iter().any(|name| *name == command) {
87        return Err(SubcommandError::unknown(command, expected));
88    }
89    Ok((command, free.collect()))
90}
91
92/// Return the default usage string used by the subcommand dispatch macros.
93pub fn usage_for_subcommand(command: &str) -> String {
94    format!("Usage: {command} [OPTIONS]")
95}
96
97/// Dispatch a `Vec<String>` of free arguments to a named subcommand.
98///
99/// Each branch is a command-name, [CommandLine] pair plus a typed handler.  The selected branch gets
100/// a concrete command-line value and the remaining free arguments.  Branch handlers return
101/// `Result<T, SubcommandError>`, which lets the macro compose recursively without enums or
102/// downcasts.
103///
104/// ```rust
105/// # use arrrg::CommandLine;
106/// # #[derive(Default, Eq, PartialEq)]
107/// # struct Top;
108/// # impl CommandLine for Top {
109/// #     fn add_opts(&self, _: Option<&str>, _: &mut getopts::Options) {}
110/// #     fn matches(&mut self, _: Option<&str>, _: &getopts::Matches) {}
111/// #     fn canonical_command_line(&self, _: Option<&str>) -> Vec<String> { vec![] }
112/// # }
113/// # #[derive(Default, Eq, PartialEq)]
114/// # struct Get;
115/// # impl CommandLine for Get {
116/// #     fn add_opts(&self, _: Option<&str>, _: &mut getopts::Options) {}
117/// #     fn matches(&mut self, _: Option<&str>, _: &getopts::Matches) {}
118/// #     fn canonical_command_line(&self, _: Option<&str>) -> Vec<String> { vec![] }
119/// # }
120/// # let (_top, free) = Top::from_arguments("Usage: tool [OPTIONS] <command>", &["get"]);
121/// let selected = arrrg::dispatch_subcommands!(free, {
122///     "get" => Get as get, get_free => {
123///         Ok((get, get_free))
124///     },
125/// });
126/// # assert!(selected.is_ok());
127/// ```
128#[macro_export]
129macro_rules! dispatch_subcommands {
130    ($free:expr, { $($name:literal => $ty:ty as $cmd:ident, $rest:ident => $body:block),+ $(,)? }) => {{
131        $crate::__dispatch_subcommands!(from_arguments, $free, {
132            $($name => $ty as $cmd, $rest => $body),+
133        })
134    }};
135}
136
137/// Dispatch subcommands while allowing non-canonical command-line argument order.
138///
139/// This is the subcommand equivalent of [CommandLine::from_arguments_relaxed].
140#[macro_export]
141macro_rules! dispatch_subcommands_relaxed {
142    ($free:expr, { $($name:literal => $ty:ty as $cmd:ident, $rest:ident => $body:block),+ $(,)? }) => {{
143        $crate::__dispatch_subcommands!(from_arguments_relaxed, $free, {
144            $($name => $ty as $cmd, $rest => $body),+
145        })
146    }};
147}
148
149#[doc(hidden)]
150#[macro_export]
151macro_rules! __dispatch_subcommands {
152    ($parser:ident, $free:expr, { $($name:literal => $ty:ty as $cmd:ident, $rest:ident => $body:block),+ $(,)? }) => {{
153        match $crate::split_subcommand($free, &[$($name),+]) {
154            Ok((__arrrg_command, __arrrg_rest)) => {
155                match __arrrg_command.as_str() {
156                    $(
157                        $name => {
158                            let __arrrg_args = __arrrg_rest
159                                .iter()
160                                .map(::std::string::String::as_str)
161                                .collect::<::std::vec::Vec<_>>();
162                            let __arrrg_usage = $crate::usage_for_subcommand($name);
163                            let ($cmd, $rest) =
164                                <$ty as $crate::CommandLine>::$parser(&__arrrg_usage, &__arrrg_args);
165                            $body
166                        }
167                    )+
168                    _ => unreachable!("split_subcommand returned an unexpected command"),
169                }
170            }
171            Err(__arrrg_err) => Err(__arrrg_err),
172        }
173    }};
174}
175
176//////////////////////////////////////////// CommandLine ///////////////////////////////////////////
177
178/// [CommandLine] creates a command line parser for anyone who implements [CommandLine::add_opts],
179/// [CommandLine::matches] and [CommandLine::canonical_command_line].  This is a wrapper around
180/// getopts to tie together options and matches.
181pub trait CommandLine: Sized + Default + Eq + PartialEq {
182    /// Add options to the getopts parser.
183    fn add_opts(&self, prefix: Option<&str>, opts: &mut getopts::Options);
184
185    /// Assign values to self using the provided getopts matches.
186    fn matches(&mut self, prefix: Option<&str>, matches: &getopts::Matches);
187
188    /// Return the canonical command line for this [CommandLine].
189    fn canonical_command_line(&self, prefix: Option<&str>) -> Vec<String>;
190
191    /// Parse from the command line.  This function will panic if a non-canonical command line is
192    /// provided.
193    fn from_command_line(usage: &str) -> (Self, Vec<String>) {
194        let args: Vec<String> = std::env::args().collect();
195        let args: Vec<&str> = args.iter().map(AsRef::as_ref).collect();
196        Self::from_arguments(usage, &args[1..])
197    }
198
199    /// Parse from the command line.  This function will allow a non-canonical command line to
200    /// execute.
201    fn from_command_line_relaxed(usage: &str) -> (Self, Vec<String>) {
202        let args: Vec<String> = std::env::args().collect();
203        let args: Vec<&str> = args.iter().map(AsRef::as_ref).collect();
204        Self::from_arguments_relaxed(usage, &args[1..])
205    }
206
207    /// Parse from the provided arguments.  This function will panic if a non-canonical command
208    /// line is provided.
209    fn from_arguments(usage: &str, args: &[&str]) -> (Self, Vec<String>) {
210        let (command_line, free) = Self::from_arguments_relaxed(usage, args);
211        let mut reconstructed_args = command_line.canonical_command_line(None);
212        let mut free_p = free.clone();
213        reconstructed_args.append(&mut free_p);
214        let mut args = args.to_vec();
215        args.retain(|a| *a != "--");
216        reconstructed_args.retain(|a| *a != "--");
217        if args != reconstructed_args {
218            panic!(
219                "non-canonical commandline specified:
220provided: {:?}
221expected: {:?}
222check argument order amongst other differences",
223                &args, reconstructed_args
224            );
225        }
226        (command_line, free)
227    }
228
229    /// Parse from the provided arguments.  This function will allow a non-canonical command line to
230    /// execute.
231    fn from_arguments_relaxed(usage: &str, args: &[&str]) -> (Self, Vec<String>) {
232        let mut command_line = Self::default();
233        let mut opts = getopts::Options::new();
234        opts.parsing_style(getopts::ParsingStyle::StopAtFirstFree);
235        opts.long_only(true);
236        opts.optflag("h", "help", "Print this help menu.");
237        opts.optflag(
238            "",
239            "completions",
240            "Print zsh completion script for this command.",
241        );
242        command_line.add_opts(None, &mut opts);
243
244        if args
245            .iter()
246            .any(|arg| arg == &"--completions" || arg == &"-completions")
247        {
248            Self::completions(
249                &mut command_line,
250                &opts,
251                usage,
252                &std::env::args()
253                    .next()
254                    .unwrap_or_else(|| "command".to_string()),
255            );
256            return (command_line, vec![]);
257        }
258
259        let matches = match opts.parse(args) {
260            Ok(matches) => matches,
261            Err(Fail::OptionMissing(which)) => {
262                Self::error(&mut command_line, format!("missing argument: --{which}"));
263                Self::usage(&mut command_line, opts, usage);
264                return (command_line, vec![]);
265            }
266            Err(err) => {
267                Self::error(
268                    &mut command_line,
269                    format!("could not parse command line: {err}"),
270                );
271                Self::exit(&mut command_line, 64);
272                return (command_line, vec![]);
273            }
274        };
275        if matches.opt_present("h") {
276            Self::usage(&mut command_line, opts, usage);
277            return (command_line, vec![]);
278        }
279        if matches.opt_present("completions") {
280            Self::completions(
281                &mut command_line,
282                &opts,
283                usage,
284                &std::env::args()
285                    .next()
286                    .unwrap_or_else(|| "command".to_string()),
287            );
288            return (command_line, vec![]);
289        }
290        command_line.matches(None, &matches);
291        let free: Vec<String> = matches.free.to_vec();
292        (command_line, free)
293    }
294
295    /// Display the usage and exit 1.
296    fn usage(&mut self, opts: getopts::Options, brief: &str) {
297        self.error(opts.usage(brief));
298        self.exit(1);
299    }
300
301    /// Emit a zsh completion script for the current command line and terminate successfully.
302    ///
303    /// The function derives the completion entries from the full option set already defined on `opts`
304    /// and prints a minimal zsh completion function scoped to the executable in `command_path`.
305    ///
306    /// # Examples
307    ///
308    /// ```rust,no_run
309    /// use arrrg::CommandLine;
310    ///
311    /// #[derive(Default, Eq, PartialEq)]
312    /// struct MyCmd;
313    ///
314    /// impl CommandLine for MyCmd {
315    ///     fn add_opts(&self, _prefix: Option<&str>, opts: &mut getopts::Options) {
316    ///         opts.optflag("h", "help", "print help");
317    ///     }
318    ///
319    ///     fn matches(&mut self, _prefix: Option<&str>, _matches: &getopts::Matches) {}
320    ///
321    ///     fn canonical_command_line(&self, _prefix: Option<&str>) -> Vec<String> {
322    ///         vec!["my-cmd".to_string()]
323    ///     }
324    /// }
325    /// ```
326    /// In normal usage, invoking `--completions` with a live implementation will print text and
327    /// exit the process with code `0`.
328    fn completions(&mut self, opts: &getopts::Options, usage: &str, command_path: &str) {
329        print!("{}", zsh_completions(opts, usage, command_path));
330        self.exit(0);
331    }
332
333    /// Report an error.
334    fn error(&mut self, msg: impl AsRef<str>) {
335        eprintln!("{}", msg.as_ref());
336    }
337
338    /// Exit with the provided status.
339    fn exit(&mut self, status: i32) {
340        std::process::exit(status);
341    }
342}
343
344fn zsh_completion_entries(opts: &getopts::Options, usage: &str) -> Vec<String> {
345    let usage = opts.usage(usage);
346    let mut options = Vec::new();
347    let mut collecting = false;
348    for line in usage.lines() {
349        if line.starts_with("Options:") {
350            collecting = true;
351            continue;
352        }
353        if !collecting {
354            continue;
355        }
356        let trimmed = line.trim();
357        if trimmed.is_empty() {
358            continue;
359        }
360        if !trimmed.starts_with('-') {
361            continue;
362        }
363        let mut split_idx = None;
364        for (i, b) in trimmed.bytes().enumerate() {
365            if b == b'\t' {
366                split_idx = Some(i);
367                break;
368            }
369            if b.is_ascii_whitespace()
370                && i + 1 < trimmed.len()
371                && trimmed.as_bytes()[i + 1].is_ascii_whitespace()
372            {
373                split_idx = Some(i);
374                break;
375            }
376        }
377        let flag_part = if let Some(i) = split_idx {
378            trimmed[..i].trim().to_string()
379        } else {
380            trimmed.to_string()
381        };
382        if flag_part.is_empty() {
383            continue;
384        }
385        let desc_part = if let Some(i) = split_idx {
386            trimmed[i..].trim()
387        } else {
388            ""
389        };
390        for name in parse_option_names(&flag_part) {
391            options.push(format_completion_option(&name, desc_part));
392        }
393    }
394    options
395}
396
397fn parse_option_names(flag_part: &str) -> Vec<String> {
398    let mut names = Vec::new();
399    for raw_name in flag_part.split(',') {
400        let mut pieces = raw_name.split_whitespace().filter(|x| !x.is_empty());
401        if let Some(name) = pieces.next()
402            && name.starts_with('-')
403        {
404            names.push(name.to_string());
405        }
406    }
407    names
408}
409
410fn format_completion_option(name: &str, desc: &str) -> String {
411    let desc = desc.replace('\'', "'\"'\"'");
412    format!("{}[{}]", name, desc)
413}
414
415fn zsh_completions(opts: &getopts::Options, usage: &str, command_path: &str) -> String {
416    let opts = zsh_completion_entries(opts, usage);
417    let command = command_path
418        .rsplit('/')
419        .next()
420        .unwrap_or(command_path)
421        .to_string();
422    let mut function: String = command
423        .chars()
424        .map(|c| if c.is_ascii_alphanumeric() { c } else { '_' })
425        .collect();
426    if function.chars().next().is_some_and(|c| c.is_ascii_digit()) {
427        function.insert(0, '_');
428    }
429
430    let mut entries = Vec::new();
431    let mut seen = HashSet::new();
432    for opt in opts {
433        if seen.insert(opt.clone()) {
434            entries.push(format!("        '{opt}'"));
435        }
436    }
437
438    let mut lines = vec![format!("#compdef {command}"), format!("_{function}() {{")];
439    if entries.is_empty() {
440        lines.push("    _arguments".to_string());
441    } else {
442        lines.push("    _arguments -s \\".to_string());
443        for (i, entry) in entries.iter().enumerate() {
444            let continuation = if i + 1 == entries.len() { "" } else { " \\" };
445            lines.push(format!("{entry}{continuation}"));
446        }
447    }
448    lines.extend_from_slice(&["}".to_string(), format!("compdef _{function} {command}")]);
449    lines.join("\n") + "\n"
450}
451
452#[cfg(test)]
453mod tests {
454    use super::*;
455
456    #[test]
457    fn zsh_completions_have_help_and_completions() {
458        let mut opts = getopts::Options::new();
459        opts.optflag("h", "help", "Print this help menu.");
460        opts.optflag(
461            "",
462            "completions",
463            "Print zsh completion script for this command.",
464        );
465        let script = zsh_completions(&opts, "Usage: test", "test-binary");
466        assert!(script.contains("#compdef test-binary"));
467        assert!(script.contains("_test_binary()"));
468        assert!(script.contains("'--help[Print this help menu.]'"));
469        assert!(script.contains("'--completions[Print zsh completion script for this command.]'"));
470    }
471
472    #[test]
473    fn zsh_completion_entries_include_short_and_long_names() {
474        let mut opts = getopts::Options::new();
475        opts.optflag("h", "help", "Print this help menu.");
476        opts.optflag("v", "verbose", "Print verbose output.");
477        let entries = zsh_completion_entries(&opts, "Usage: test");
478
479        assert!(entries.contains(&"-h[Print this help menu.]".to_string()));
480        assert!(entries.contains(&"--help[Print this help menu.]".to_string()));
481        assert!(entries.contains(&"-v[Print verbose output.]".to_string()));
482        assert!(entries.contains(&"--verbose[Print verbose output.]".to_string()));
483    }
484
485    #[test]
486    fn zsh_completion_entries_escapes_single_quotes() {
487        let mut opts = getopts::Options::new();
488        opts.optflag("", "with-space", "It'll parse correctly.");
489        let entries = zsh_completion_entries(&opts, "Usage: test");
490
491        assert!(entries.contains(&"--with-space[It'\"'\"'ll parse correctly.]".to_string()));
492    }
493
494    #[test]
495    fn zsh_completions_normalize_function_name() {
496        let mut opts = getopts::Options::new();
497        opts.optflag("h", "help", "Print this help menu.");
498        let script = zsh_completions(&opts, "Usage: test", "./bin/my-cmd");
499
500        assert!(script.contains("#compdef my-cmd"));
501        assert!(script.contains("_my_cmd()"));
502        assert!(script.contains("compdef _my_cmd my-cmd"));
503    }
504
505    #[derive(Default, Eq, PartialEq)]
506    struct TestCommandLine {
507        required: String,
508    }
509
510    impl CommandLine for TestCommandLine {
511        fn add_opts(&self, _prefix: Option<&str>, opts: &mut getopts::Options) {
512            opts.reqopt("", "chooser-mode", "required mode", "METHOD");
513        }
514
515        fn matches(&mut self, _prefix: Option<&str>, matches: &getopts::Matches) {
516            if let Some(mode) = matches.opt_str("chooser-mode") {
517                self.required = mode;
518            }
519        }
520
521        fn canonical_command_line(&self, _prefix: Option<&str>) -> Vec<String> {
522            let mut result = vec!["test".to_string()];
523            if !self.required.is_empty() {
524                result.push("--chooser-mode".to_string());
525                result.push(self.required.to_string());
526            }
527            result
528        }
529    }
530
531    #[test]
532    fn completions_bypass_missing_required_args() {
533        let (command_line, _) = NoExitCommandLine::<TestCommandLine>::from_arguments_relaxed(
534            "Usage: test",
535            &["--completions"],
536        );
537        let (_, _, status) = command_line.into_parts();
538        assert_eq!(status, 0);
539    }
540}
541
542///////////////////////////////////////// NoExitCommandLine ////////////////////////////////////////
543
544/// A non-exiting wrapper for command line parsing.  Will store command line in 0, messages in
545/// element 1, exit status in 2.
546#[derive(Default, Eq, PartialEq)]
547pub struct NoExitCommandLine<T: CommandLine>(T, Vec<String>, i32);
548
549impl<T: CommandLine> AsRef<T> for NoExitCommandLine<T> {
550    fn as_ref(&self) -> &T {
551        &self.0
552    }
553}
554
555impl<T: CommandLine> NoExitCommandLine<T> {
556    pub fn into_inner(self) -> T {
557        self.0
558    }
559
560    pub fn into_parts(self) -> (T, Vec<String>, i32) {
561        (self.0, self.1, self.2)
562    }
563}
564
565impl<T: CommandLine> CommandLine for NoExitCommandLine<T> {
566    fn add_opts(&self, prefix: Option<&str>, opts: &mut getopts::Options) {
567        self.0.add_opts(prefix, opts);
568    }
569
570    fn matches(&mut self, prefix: Option<&str>, matches: &getopts::Matches) {
571        self.0.matches(prefix, matches);
572    }
573
574    fn canonical_command_line(&self, prefix: Option<&str>) -> Vec<String> {
575        self.0.canonical_command_line(prefix)
576    }
577
578    fn error(&mut self, msg: impl AsRef<str>) {
579        self.1.push(msg.as_ref().to_string());
580    }
581
582    fn exit(&mut self, status: i32) {
583        self.2 = status;
584    }
585}
586
587//////////////////////////////////////////// macro utils ///////////////////////////////////////////
588
589#[doc(hidden)]
590pub fn getopt_str(prefix: Option<&str>, field_arg: &str) -> String {
591    match prefix {
592        Some(prefix) => {
593            format!("{prefix}-{field_arg}")
594        }
595        None => field_arg.to_string(),
596    }
597}
598
599#[doc(hidden)]
600pub fn dashed_str(prefix: Option<&str>, field_arg: &str) -> String {
601    format!("--{}", getopt_str(prefix, field_arg))
602}
603
604#[doc(hidden)]
605pub fn parse_field<T>(arg_str: &str, s: &str) -> T
606where
607    T: FromStr,
608    <T as FromStr>::Err: std::fmt::Display,
609{
610    match s.parse::<T>() {
611        Ok(t) => t,
612        Err(err) => {
613            panic!("field --{arg_str} is unparseable: {err}");
614        }
615    }
616}
arrrg/lib.rs

arrrg/
lib.rs
