Skip to main content

ashlar_cli/
args.rs

1//! A tiny argument parser: flags (`--json`), options (`--connect 7981`, `--connect=7981`) and
2//! positionals, taken in any order. Whatever a verb does not take is an error at [`Args::finish`].
3
4/// The arguments after the verb.
5#[derive(Debug, Clone, PartialEq, Default)]
6pub struct Args {
7    items: Vec<String>,
8    /// Items after a bare `--` are positionals even when they start with `--`.
9    rest: Vec<String>,
10}
11
12impl Args {
13    pub fn new(items: impl IntoIterator<Item = impl Into<String>>) -> Self {
14        let mut all: Vec<String> = items.into_iter().map(Into::into).collect();
15        let rest = match all.iter().position(|a| a == "--") {
16            Some(i) => {
17                let tail = all.split_off(i);
18                tail.into_iter().skip(1).collect()
19            }
20            None => Vec::new(),
21        };
22        Self { items: all, rest }
23    }
24
25    /// Take `--name` if present.
26    pub fn flag(&mut self, name: &str) -> bool {
27        match self.items.iter().position(|a| a == name) {
28            Some(i) => {
29                self.items.remove(i);
30                true
31            }
32            None => false,
33        }
34    }
35
36    /// Take `--name VALUE` or `--name=VALUE`.
37    pub fn value(&mut self, name: &str) -> Result<Option<String>, String> {
38        let prefix = format!("{name}=");
39        if let Some(i) = self.items.iter().position(|a| a.starts_with(&prefix)) {
40            let a = self.items.remove(i);
41            return Ok(Some(a.get(prefix.len()..).unwrap_or("").to_string()));
42        }
43        let Some(i) = self.items.iter().position(|a| a == name) else { return Ok(None) };
44        self.items.remove(i);
45        if i < self.items.len() { Ok(Some(self.items.remove(i))) } else { Err(format!("{name} needs a value")) }
46    }
47
48    /// Take `--name N` as a number.
49    pub fn number<T: std::str::FromStr>(&mut self, name: &str) -> Result<Option<T>, String> {
50        match self.value(name)? {
51            Some(v) => v.parse().map(Some).map_err(|_| format!("{name}: `{v}` is not a valid number")),
52            None => Ok(None),
53        }
54    }
55
56    /// Take the next positional (`-` counts as one: stdin). Take options first, so an option's
57    /// value is not mistaken for a positional.
58    pub fn positional(&mut self) -> Option<String> {
59        match self.items.iter().position(|a| !a.starts_with("--")) {
60            Some(i) => Some(self.items.remove(i)),
61            None => (!self.rest.is_empty()).then(|| self.rest.remove(0)),
62        }
63    }
64
65    /// Whatever was not taken, in the original order (a bare `--` before the items that followed
66    /// it), for a verb that parses its arguments itself.
67    pub fn into_raw(self) -> Vec<String> {
68        let mut v = self.items;
69        if !self.rest.is_empty() {
70            v.push("--".into());
71            v.extend(self.rest);
72        }
73        v
74    }
75
76    /// Everything was taken; otherwise the first leftover is an error.
77    pub fn finish(self) -> Result<(), String> {
78        match self.items.first().or(self.rest.first()) {
79            None => Ok(()),
80            Some(a) if a.starts_with("--") => Err(format!("unknown option {a}")),
81            Some(a) => Err(format!("unexpected argument `{a}`")),
82        }
83    }
84}
85
86#[cfg(test)]
87mod tests {
88    use super::*;
89
90    #[test]
91    fn the_raw_arguments_come_back_in_order() {
92        let mut a = Args::new(["x", "--json", "y", "--", "--z"]);
93        assert!(a.flag("--json"));
94        assert_eq!(a.into_raw(), ["x", "y", "--", "--z"]);
95    }
96
97    #[test]
98    fn flags_values_and_positionals_in_any_order() {
99        let mut a = Args::new(["--json", "doc.add", "--connect=7981", "{\"x\":1}", "--depth", "3"]);
100        assert!(a.flag("--json"));
101        assert!(!a.flag("--json"));
102        assert_eq!(a.number::<u16>("--connect"), Ok(Some(7981)));
103        assert_eq!(a.number::<u32>("--depth"), Ok(Some(3)));
104        assert_eq!(a.positional().as_deref(), Some("doc.add"));
105        assert_eq!(a.positional().as_deref(), Some("{\"x\":1}"));
106        assert_eq!(a.positional(), None);
107        assert_eq!(a.finish(), Ok(()));
108    }
109
110    #[test]
111    fn leftovers_and_missing_values_are_errors() {
112        assert_eq!(Args::new(["--nope"]).finish(), Err("unknown option --nope".into()));
113        assert_eq!(Args::new(["extra"]).finish(), Err("unexpected argument `extra`".into()));
114        assert_eq!(Args::new(["--connect"]).value("--connect"), Err("--connect needs a value".into()));
115        assert!(Args::new(["--connect", "x"]).number::<u16>("--connect").is_err());
116        assert!(Args::new(["--connect", "70000"]).number::<u16>("--connect").is_err());
117    }
118
119    #[test]
120    fn double_dash_ends_options() {
121        let mut a = Args::new(["exec", "--", "--weird-id"]);
122        assert_eq!(a.positional().as_deref(), Some("exec"));
123        assert_eq!(a.positional().as_deref(), Some("--weird-id"));
124        assert_eq!(a.finish(), Ok(()));
125        let mut a = Args::new(["-"]);
126        assert_eq!(a.positional().as_deref(), Some("-"));
127    }
128}