Skip to main content

az_rs/
arg.rs

1use anyhow::Result;
2
3#[derive(Debug, Clone, PartialEq)]
4pub enum Arg {
5    Optional(String, Option<String>), // Can be: --enable, --enable=true, --foo bar
6    Positional(String),
7}
8
9#[derive(Debug, Clone)]
10pub struct CliInput {
11    pub args: Vec<Arg>,
12}
13
14impl CliInput {
15    pub fn new<I, S>(args: I) -> Result<Self>
16    where
17        I: IntoIterator<Item = S>,
18        S: AsRef<str>,
19    {
20        let mut result = Vec::new();
21        let mut see_opt = false;
22        let mut args = args.into_iter().peekable();
23
24        while let Some(arg) = args.next() {
25            if let Some(arg) = arg.as_ref().strip_prefix("-") {
26                see_opt = true;
27                let mut arg = arg;
28                if let Some(arg2) = arg.strip_prefix("-") {
29                    arg = arg2;
30                }
31                if let Some(eq_idx) = arg.find('=') {
32                    // Handle --key=value
33                    let key = arg[..eq_idx].to_string();
34                    let value = Some(arg[eq_idx + 1..].to_string());
35                    result.push(Arg::Optional(key, value));
36                } else {
37                    // Handle --key [value] or just --flag
38                    let key = arg.to_string();
39                    match args.peek() {
40                        Some(next) if !next.as_ref().starts_with("-") => {
41                            let value = args.next().map(|v| String::from(v.as_ref()));
42                            result.push(Arg::Optional(key, value));
43                        }
44                        _ => {
45                            // It's a flag with no value
46                            result.push(Arg::Optional(key, None));
47                        }
48                    }
49                }
50            } else {
51                if see_opt {
52                    anyhow::bail!("optional raw arguments must follow positional raw arguments");
53                }
54                // Positional argument
55                result.push(Arg::Positional(String::from(arg.as_ref())));
56            }
57        }
58        Ok(Self { args: result })
59    }
60
61    pub fn is_help(&self) -> bool {
62        self.args
63            .iter()
64            .filter(|arg| {
65                if let Arg::Optional(k, _) = arg {
66                    return k == "h" || k == "help";
67                }
68                false
69            })
70            .count()
71            != 0
72    }
73
74    pub fn pos_args(&self) -> Vec<&str> {
75        self.args
76            .iter()
77            .filter_map(|arg| {
78                if let Arg::Positional(arg) = arg {
79                    Some(arg.as_str())
80                } else {
81                    None
82                }
83            })
84            .collect()
85    }
86
87    pub fn len(&self) -> usize {
88        self.args.len()
89    }
90
91    pub fn is_empty(&self) -> bool {
92        self.len() == 0
93    }
94}
95
96#[cfg(test)]
97mod test {
98    use crate::arg::{Arg, CliInput};
99
100    #[test]
101    fn new_cli_input() {
102        assert_eq!(
103            CliInput::new(vec!["foo"]).unwrap().args,
104            vec![Arg::Positional(String::from("foo"))]
105        );
106        assert_eq!(
107            CliInput::new(vec!["foo", "bar"]).unwrap().args,
108            vec![
109                Arg::Positional(String::from("foo")),
110                Arg::Positional(String::from("bar")),
111            ]
112        );
113        assert_eq!(
114            CliInput::new(vec!["foo", "--bar"]).unwrap().args,
115            vec![
116                Arg::Positional(String::from("foo")),
117                Arg::Optional(String::from("bar"), None),
118            ]
119        );
120        assert_eq!(
121            CliInput::new(vec!["foo", "-b"]).unwrap().args,
122            vec![
123                Arg::Positional(String::from("foo")),
124                Arg::Optional(String::from("b"), None),
125            ]
126        );
127        assert_eq!(
128            CliInput::new(vec!["foo", "--bar=baz"]).unwrap().args,
129            vec![
130                Arg::Positional(String::from("foo")),
131                Arg::Optional(String::from("bar"), Some(String::from("baz"))),
132            ]
133        );
134        assert_eq!(
135            CliInput::new(vec!["foo", "-b=baz"]).unwrap().args,
136            vec![
137                Arg::Positional(String::from("foo")),
138                Arg::Optional(String::from("b"), Some(String::from("baz"))),
139            ]
140        );
141        assert_eq!(
142            CliInput::new(vec!["foo", "--bar", "baz"]).unwrap().args,
143            vec![
144                Arg::Positional(String::from("foo")),
145                Arg::Optional(String::from("bar"), Some(String::from("baz"))),
146            ]
147        );
148        assert_eq!(
149            CliInput::new(vec!["foo", "-b", "baz"]).unwrap().args,
150            vec![
151                Arg::Positional(String::from("foo")),
152                Arg::Optional(String::from("b"), Some(String::from("baz"))),
153            ]
154        );
155    }
156}