core/cmd/
arg.rs

1use std::fmt;
2
3use anyhow::bail;
4
5pub(crate) trait Args: Default {
6    fn help() -> &'static str;
7
8    fn parse(input: impl AsRef<str>) -> anyhow::Result<Self>;
9
10    fn parse_inner(
11        input: impl AsRef<str>,
12        predicate: impl Fn(&mut Self, &str, Option<&ArgValue>) -> bool,
13    ) -> anyhow::Result<Self>
14    where
15        Self: Sized,
16    {
17        let mut args = Self::default();
18
19        let parsed_args: Result<Vec<_>, _> =
20            input.as_ref().split_whitespace().map(Arg::parse).collect();
21
22        for arg in parsed_args?.into_iter() {
23            if !predicate(&mut args, &arg.name, arg.value.as_ref()) {
24                bail!("unrecognized or ill-formed argument: {arg}")
25            }
26        }
27
28        Ok(args)
29    }
30}
31
32pub(crate) struct Arg {
33    name: String,
34    value: Option<ArgValue>,
35}
36
37impl Arg {
38    fn parse(input: impl AsRef<str>) -> anyhow::Result<Self> {
39        let input = input.as_ref();
40
41        let (sign, kv) = (
42            {
43                let mut iter = input.chars();
44                iter.next()
45                    .filter(|&ch| ch == '+' || ch == '-')
46                    .zip(Some(iter.as_str()))
47            },
48            input.split_once('='),
49        );
50
51        let arg = match (sign, kv) {
52            (None, Some((name, value))) => Self {
53                name: name.into(),
54                value: Some(ArgValue::KV(value.into())),
55            },
56            (Some((ch, name)), None) => Self {
57                name: name.into(),
58                value: Some(ArgValue::Bool(ch == '+')),
59            },
60            (Some(_), Some(_)) => {
61                bail!("+/-option with =value is not supported yet");
62            }
63            (None, None) => Self {
64                name: input.into(),
65                value: None,
66            },
67        };
68
69        Ok(arg)
70    }
71}
72
73impl fmt::Display for Arg {
74    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
75        match &self.value {
76            Some(ArgValue::Bool(enable)) => {
77                write!(f, "{}{}", *enable, self.name)
78            }
79            Some(ArgValue::KV(value)) => write!(f, "{}={}", self.name, value),
80            None => write!(f, "{}", self.name),
81        }
82    }
83}
84
85pub(crate) enum ArgValue {
86    Bool(bool), // `-opt`, `+opt`
87    KV(String), // `arg=abc`
88}
89
90#[macro_export]
91macro_rules! define_cmd_args {
92    ( $help:literal $(#[$attrs:meta])* $vis:vis struct $name:ident { $($body:tt)* } ) => {
93        $(#[$attrs])*
94        $vis struct $name { $($body)* }
95
96        impl $crate::cmd::Args for $name {
97            fn help() -> &'static str {
98                $help
99            }
100
101            fn parse(input: impl AsRef<str>) -> anyhow::Result<Self> {
102                Self::parse_inner(input, |args, name, value| {
103                    define_cmd_args!(@ARM, (name, value), args, $($body)*)
104                })
105            }
106        }
107    };
108    ( @ARM, $input:expr, $result:expr,
109      $(#[$attrs:meta])* $vis:vis $name:ident : bool, $($body:tt)*) => {
110        if let (stringify!($name), None) = $input {
111            $result.$name = true;
112            return true;
113        } else {
114            define_cmd_args!(@ARM, $input, $result, $($body)*)
115        }
116    };
117    ( @ARM, $input:expr, $result:expr,
118      $(#[$attrs:meta])* $vis:vis $name:ident : Option<bool>, $($body:tt)*) => {
119        if let (stringify!($name), Some($crate::cmd::ArgValue::Bool(enable))) = $input {
120            $result.$name = Some(*enable);
121            return true;
122        } else {
123            define_cmd_args!(@ARM, $input, $result, $($body)*)
124        }
125    };
126    ( @ARM, $input:expr, $result:expr,
127      $(#[$attrs:meta])* $vis:vis $name:ident : Option<String>, $($body:tt)*) => {
128        if let (stringify!($name), Some(ArgValue::KV(value))) = $input {
129            $result.$name = Some(value.into());
130            return true;
131        } else {
132            define_cmd_args!(@ARM, $input, $result, $($body)*)
133        }
134    };
135    ( @ARM, $input:expr, $result:expr,) => {
136      false
137    };
138}
139pub use define_cmd_args;
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144
145    define_cmd_args! {
146        "help text"
147
148        #[derive(PartialEq, Eq, Debug, Default)]
149        pub struct TestArgs {
150            pub help: bool,
151            opt_bool: Option<bool>,
152            opt_string: Option<String>,
153        }
154    }
155
156    #[test]
157    fn validation() {
158        assert_eq!(TestArgs::help(), "help text");
159        assert_eq!(
160            TestArgs::parse("").unwrap(),
161            TestArgs {
162                help: false,
163                opt_bool: None,
164                opt_string: None,
165            }
166        );
167
168        assert_eq!(
169            TestArgs::parse("help").unwrap(),
170            TestArgs {
171                help: true,
172                opt_bool: None,
173                opt_string: None,
174            }
175        );
176        assert!(TestArgs::parse("+help").is_err());
177        assert!(TestArgs::parse("-help").is_err());
178        assert!(TestArgs::parse("help=abc").is_err());
179
180        assert_eq!(
181            TestArgs::parse("+opt_bool").unwrap(),
182            TestArgs {
183                help: false,
184                opt_bool: Some(true),
185                opt_string: None,
186            }
187        );
188        assert_eq!(
189            TestArgs::parse("-opt_bool").unwrap(),
190            TestArgs {
191                help: false,
192                opt_bool: Some(false),
193                opt_string: None,
194            }
195        );
196        assert!(TestArgs::parse("opt_bool").is_err());
197        assert!(TestArgs::parse("opt_bool=abc").is_err());
198
199        assert_eq!(
200            TestArgs::parse("opt_string=abc").unwrap(),
201            TestArgs {
202                help: false,
203                opt_bool: None,
204                opt_string: Some("abc".into()),
205            }
206        );
207        assert!(TestArgs::parse("opt_string").is_err());
208        assert!(TestArgs::parse("+opt_string").is_err());
209        assert!(TestArgs::parse("-opt_string").is_err());
210    }
211}