arg_parser/
lib.rs

1use std::borrow::Borrow;
2use std::cell::{RefCell, RefMut};
3use std::collections::HashMap;
4use std::hash::{Hash, Hasher};
5use std::rc::Rc;
6use std::time::{SystemTime, UNIX_EPOCH};
7
8#[derive(Clone, Debug, Eq, PartialEq)]
9/// The parameter styles for short, e.g. `-s`,
10/// and for long, e.g. `--long`
11pub enum Param {
12    Short(char),
13    Long(String),
14}
15
16impl Borrow<str> for Param {
17    fn borrow(&self) -> &str {
18        if let Param::Long(ref string) = *self {
19            string
20        } else {
21            ""
22        }
23    }
24}
25
26impl Borrow<char> for Param {
27    fn borrow(&self) -> &char {
28        if let Param::Short(ref ch) = *self {
29            ch
30        } else {
31            const CH: &'static char = &'\0';
32            CH
33        }
34    }
35}
36
37impl Hash for Param {
38    fn hash<H: Hasher>(&self, state: &mut H) {
39        match *self {
40            Param::Short(ref c) => c.hash(state),
41            Param::Long(ref s) => s.hash(state),
42        }
43    }
44}
45
46#[derive(Clone, Debug, Default, Eq, PartialEq)]
47/// The Right Hand Side type
48struct Rhs<T> {
49    /// The RHS value
50    value: T,
51    /// Counts the number of times a flag/opt has been in use on the command
52    occurrences: usize,
53}
54
55impl<T> Rhs<T> {
56    fn new(value: T) -> Self {
57        Rhs {
58            value: value,
59            occurrences: 0,
60        }
61    }
62}
63
64#[derive(Clone, Debug, Eq, PartialEq)]
65/// The Value for each parameter
66enum Value {
67    Flag(Rhs<Rc<RefCell<bool>>>),
68    /// The RHS String value is shared between both short and long parameters
69    Opt {
70        rhs: Rhs<Rc<RefCell<String>>>,
71        found: Rc<RefCell<bool>>
72    },
73    Setting {
74        rhs: Rhs<Rc<RefCell<String>>>,
75        found: Rc<RefCell<bool>>
76    },
77}
78
79impl Value {
80    fn new_opt(value: Rc<RefCell<String>>, found: Rc<RefCell<bool>>) -> Self {
81        Value::Opt {
82            rhs: Rhs::new(value),
83            found,
84        }
85    }
86
87    fn new_setting(value: Rc<RefCell<String>>, found: Rc<RefCell<bool>>) -> Self {
88        Value::Setting {
89            rhs: Rhs::new(value),
90            found,
91        }
92    }
93}
94
95/// Our homebrewed Arg Parser
96#[derive(Clone, Debug, Default)]
97pub struct ArgParser {
98    params: HashMap<Param, Value>,
99    invalid: Vec<Param>,
100    garbage: (RefCell<bool>, RefCell<String>),
101    pub args: Vec<String>,
102}
103
104impl ArgParser {
105    /// Create a new ArgParser object
106    ///
107    /// The capacity specifies the initial capacity of the number parameters.
108    /// Always good to set it at the number of flags and opts total.
109    /// (Used for initial hashmap allocation size)
110    pub fn new(capacity: usize) -> Self {
111        ArgParser {
112            params: HashMap::with_capacity(capacity),
113            invalid: Vec::new(),
114            garbage: (RefCell::new(false), RefCell::new(String::with_capacity(0))),
115            args: Vec::new(),
116        }
117    }
118
119    /// Builder method for adding both short and long flags
120    ///
121    /// Flags are just parameters that have no assigned values. They are used
122    /// for when certain features or options have been enabled for the application
123    ///
124    /// # Examples
125    /// ```text
126    /// > ls -l --human-readable
127    ///   ^  ^  ^
128    ///   |  |  |
129    ///   |  |  `-- A long flag to enable human readable numbers.
130    ///   |  `-- A short flag to enable the long format.
131    ///   `-- The command to list files.
132    /// ```
133    pub fn add_flag(mut self, flags: &[&str]) -> Self {
134        let value = Rc::new(RefCell::new(bool::default()));
135        for flag in flags.iter() {
136            if flag.len() == 1 {
137                if let Some(short) = flag.chars().next() {
138                    self.params.insert(Param::Short(short), Value::Flag(Rhs::new(value.clone())));
139                }
140            } else if !flag.is_empty() {
141                self.params.insert(Param::Long((*flag).to_owned()), Value::Flag(Rhs::new(value.clone())));
142            }
143        }
144        self
145    }
146
147    /// Builder method for adding both short and long opts
148    ///
149    /// Opts are parameters that hold assigned values. They are used
150    /// for when certain features or options have been enabled for the application
151    ///
152    /// # Examples
153    /// ```text
154    /// > ls -T 4 --color=always
155    ///   ^  ^    ^
156    ///   |  |    |
157    ///   |  |    `-- A long opt to enable the use of color with value `always`.
158    ///   |  `-- A short opt to set tab size to the value `4`.
159    ///   `-- The command to list files.
160    /// ```
161    pub fn add_opt(mut self, short: &str, long: &str) -> Self {
162        let value = Rc::new(RefCell::new("".to_owned()));
163        let found = Rc::new(RefCell::new(false));
164        if let Some(short) = short.chars().next() {
165            self.params.insert(Param::Short(short), Value::new_opt(value.clone(), found.clone()));
166        }
167        if !long.is_empty() {
168            self.params.insert(Param::Long(long.to_owned()), Value::new_opt(value, found));
169        }
170        self
171    }
172    
173    /// Add an opt with a default value. Works the same way as [add_opt](struct.ArgParser.html#method.add_opt).
174    pub fn add_opt_default(mut self, short: &str, long: &str, default: &str) -> Self {
175        let value = Rc::new(RefCell::new(default.to_owned()));
176        let found = Rc::new(RefCell::new(true));
177        if let Some(short) = short.chars().next() {
178            self.params.insert(Param::Short(short), Value::new_opt(value.clone(), found.clone()));
179        }
180        if !long.is_empty() {
181            self.params.insert(Param::Long(long.to_owned()), Value::new_opt(value, found));
182        }
183        self
184    }
185
186    /// Builder method for adding settings
187    ///
188    /// Settings are parameters that hold assigned values. They are used
189    /// in some applications such as dd
190    ///
191    /// # Examples
192    /// ```text
193    /// > dd if=/path/file
194    ///   ^  ^
195    ///   |  |
196    ///   |  |
197    ///   |  `-- The setting set to /path/file
198    ///   `-- The command to list files.
199    /// ```
200    pub fn add_setting(mut self, setting: &str) -> Self {
201        let value = Rc::new(RefCell::new("".to_owned()));
202        let found = Rc::new(RefCell::new(false));
203        if !setting.is_empty() {
204            self.params.insert(Param::Long(setting.to_owned()), Value::new_setting(value, found));
205        }
206        self
207    }
208
209    /// Add a setting with a default value. Works the same way as [add_setting](struct.ArgParser.html#method.add_setting)
210    pub fn add_setting_default(mut self, setting: &str, default: &str) -> Self {
211        let value = Rc::new(RefCell::new(default.to_owned()));
212        let found = Rc::new(RefCell::new(true));
213        if !setting.is_empty() {
214            self.params.insert(Param::Long(setting.to_owned()), Value::new_setting(value, found));
215        }
216        self
217    }
218
219    /// Start parsing user inputted args for which flags and opts are used at
220    /// runtime. The rest of the args that are not associated to opts get added
221    /// to `ArgParser.args`.
222    pub fn parse<A: Iterator<Item = String>>(&mut self, args: A) {
223        let mut args = args.skip(1);
224        while let Some(arg) = args.next() {
225            if arg.starts_with("--") {
226                // Remove both dashes
227                let arg = &arg[2..];
228                if arg.is_empty() {
229                    //Arg `--` means we are done parsing args, collect the rest
230                    self.args.extend(args);
231                    break;
232                }
233                if let Some(i) = arg.find('=') {
234                    let (lhs, rhs) = arg.split_at(i);
235                    let rhs = &rhs[1..]; // slice off the `=` char
236                    match self.params.get_mut(lhs) {
237                        Some(&mut Value::Opt { rhs: ref mut opt_rhs, ref mut found }) => {
238                            if (*opt_rhs.value).borrow().is_empty() {
239                                opt_rhs.occurrences = 1;
240                            } else {
241                                opt_rhs.occurrences += 1;
242                            }
243                            (*opt_rhs.value).borrow_mut().clear();
244                            (*opt_rhs.value).borrow_mut().push_str(rhs);
245                            *(*found).borrow_mut() = true;
246                        }
247                        _ => self.invalid.push(Param::Long(lhs.to_owned())),
248                    }
249                } else {
250                    match self.params.get_mut(arg) {
251                        Some(&mut Value::Flag(ref mut rhs)) => {
252                            *(*rhs.value).borrow_mut() = true;
253                            rhs.occurrences += 1;
254                        }
255                        Some(&mut Value::Opt { ref mut rhs, ref mut found }) => {
256                            rhs.occurrences += 1;
257                            *(*found).borrow_mut() = true;
258                        }
259                        _ => self.invalid.push(Param::Long(arg.to_owned())),
260                    }
261                }
262            } else if arg.starts_with("-") && arg != "-" {
263                let mut chars = arg[1..].chars();
264                while let Some(ch) = chars.next() {
265                    match self.params.get_mut(&ch) {
266                        Some(&mut Value::Flag(ref mut rhs)) => {
267                            *(*rhs.value).borrow_mut() = true;
268                            rhs.occurrences += 1;
269                        }
270                        Some(&mut Value::Opt { ref mut rhs, ref mut found }) => {
271                            let rest: String = chars.collect();
272                            if !rest.is_empty() {
273                                *(*rhs.value).borrow_mut() = rest;
274                                *(*found).borrow_mut() = true;
275                            } else {
276                                *(*rhs.value).borrow_mut() = args.next()
277                                    .map(|a| {
278                                             *(*found).borrow_mut() = true;
279                                             a
280                                         })
281                                    .unwrap_or("".to_owned());
282                            }
283                            break;
284                        }
285                        Some(&mut Value::Setting { .. }) => self.invalid.push(Param::Short(ch)),
286                        None => self.invalid.push(Param::Short(ch)),
287                    }
288                }
289            } else if arg.contains("=") {
290                if arg.is_empty() {
291                    //Arg `--` means we are done parsing args, collect the rest
292                    self.args.extend(args);
293                    break;
294                }
295                if let Some(i) = arg.find('=') {
296                    let (lhs, rhs) = arg.split_at(i);
297                    let rhs = &rhs[1..]; // slice off the `=` char
298                    match self.params.get_mut(lhs) {
299                        Some(&mut Value::Setting { rhs: ref mut opt_rhs, ref mut found }) => {
300                            if (*opt_rhs.value).borrow().is_empty() {
301                                opt_rhs.occurrences = 1;
302                            } else {
303                                opt_rhs.occurrences += 1;
304                            }
305                            (*opt_rhs.value).borrow_mut().clear();
306                            (*opt_rhs.value).borrow_mut().push_str(rhs);
307                            *(*found).borrow_mut() = true;
308                        }
309                        _ => self.invalid.push(Param::Long(lhs.to_owned())),
310                    }
311                }
312            } else {
313                self.args.push(arg);
314            }
315        }
316    }
317
318    /// Get the number of times a flag or opt has been found after parsing.
319    pub fn count<P: Hash + Eq + ?Sized>(&self, name: &P) -> usize
320        where Param: Borrow<P>
321    {
322        match self.params.get(name) {
323            Some(&Value::Flag(ref rhs)) => rhs.occurrences,
324            Some(&Value::Opt { ref rhs, .. }) => rhs.occurrences,
325            _ => 0,
326        }
327    }
328
329    /// Check if a flag or opt has been found after initialization.
330    /// While the short form of the opt or flag will be recognized,
331    /// the long form should be used to enhance code readability.
332    ///
333    /// # Examples
334    ///
335    /// ```
336    /// if parser.found("my-long-flag") {
337    ///     //Do things...
338    /// }
339    /// ```
340    pub fn found<P: Hash + Eq + ?Sized>(&self, name: &P) -> bool
341        where Param: Borrow<P>
342    {
343        match self.params.get(name) {
344            Some(&Value::Flag(ref rhs)) => *(*rhs.value).borrow_mut(),
345            Some(&Value::Opt { ref found, .. }) => *(**found).borrow(),
346            Some(&Value::Setting { ref found, .. }) => *(**found).borrow(),
347            _ => false,
348        }
349    }
350
351    /// Modify the state of a flag. Use `true` if the flag is to be enabled. Use `false` to
352    /// disable its use.
353    pub fn flag<F: Hash + Eq + ?Sized>(&mut self, flag: &F) -> RefMut<bool>
354        where Param: Borrow<F>
355    {
356        if let Some(&mut Value::Flag(ref mut rhs)) = self.params.get_mut(flag) {
357            return (*rhs.value).borrow_mut();
358        }
359        self.garbage.0.borrow_mut()
360    }
361
362    /// Modify the state value of an opt. Use `Some(String)` to set if the opt is to be enabled and
363    /// has been assigned a value from `String`. Use `None` to disable the opt's use.
364    pub fn opt<O: Hash + Eq + ?Sized>(&mut self, opt: &O) -> RefMut<String>
365        where Param: Borrow<O>
366    {
367        if let Some(&mut Value::Opt { ref mut rhs, .. }) = self.params.get_mut(opt) {
368            return (*rhs.value).borrow_mut();
369        }
370        self.garbage.1.borrow_mut()
371    }
372
373    /// Get the value of an Opt. If it has been set or defaulted, it will return a `Some(String)`
374    /// value otherwise it will return None.
375    pub fn get_opt<O: Hash + Eq + ?Sized>(&self, opt: &O) -> Option<String>
376        where Param: Borrow<O>
377    {
378        if let Some(&Value::Opt { ref rhs, ref found }) = self.params.get(opt) {
379            if *(**found).borrow() {
380                return Some((*rhs.value).borrow().clone());
381            }
382        }
383        None
384    }
385
386    /// Get the value of an Setting. If it has been set or defaulted, it will return a `Some(String)`
387    /// value otherwise it will return None.
388    pub fn get_setting<O: Hash + Eq + ?Sized>(&self, setting: &O) -> Option<String>
389        where Param: Borrow<O>
390    {
391        if let Some(&Value::Setting { ref rhs, ref found }) = self.params.get(setting) {
392            if *(**found).borrow() {
393                return Some((*rhs.value).borrow().clone());
394            }
395        }
396        None
397    }
398
399    pub fn found_invalid(&self) -> Result<(), String> {
400        if self.invalid.is_empty() {
401            return Ok(());
402        }
403
404        let mut and: bool = false;
405        let mut output = if self.invalid.len() == 1 {
406                "Invalid parameter"
407            } else {
408                and = true;
409                "Invalid parameters"
410            }
411            .to_owned();
412
413        let mut iter = self.invalid.iter().peekable();
414        while let Some(param) = iter.next() {
415            match param {
416                &Param::Short(ch) => {
417                    output += " '-";
418                    output.push(ch);
419                    output.push('\'');
420                }
421                &Param::Long(ref s) => {
422                    output += " '--";
423                    output += s;
424                    output.push('\'');
425                }
426            }
427            if and && iter.peek().is_some() {
428                output += " and";
429            }
430        }
431        output.push('\n');
432        Err(output)
433    }
434}
435
436pub fn format_system_time(time: SystemTime) -> String {
437    let tz_offset = 0; //TODO Apply timezone offset
438    match time.duration_since(UNIX_EPOCH) {
439        Ok(duration) => format_time(duration.as_secs() as i64, tz_offset),
440        Err(_) => "duration since epoch err".to_string(),
441    }
442}
443
444// Sweet algorithm from http://ptspts.blogspot.com/2009/11/how-to-convert-unix-timestamp-to-civil.html
445// TODO: Apply timezone offset
446pub fn get_time_tuple(mut ts: i64, tz_offset: i64) -> (i64, i64, i64, i64, i64, i64) {
447    ts += tz_offset * 3600;
448    let s = ts % 86400;
449    ts /= 86400;
450    let h = s / 3600;
451    let m = s / 60 % 60;
452    let s = s % 60;
453    let x = (ts * 4 + 102032) / 146097 + 15;
454    let b = ts + 2442113 + x - (x / 4);
455    let mut c = (b * 20 - 2442) / 7305;
456    let d = b - 365 * c - c / 4;
457    let mut e = d * 1000 / 30601;
458    let f = d - e * 30 - e * 601 / 1000;
459    if e < 14 {
460        c -= 4716;
461        e -= 1;
462    } else {
463        c -= 4715;
464        e -= 13;
465    }
466    (c, e, f, h, m, s)
467}
468
469pub fn format_time(ts: i64, tz_offset: i64) -> String {
470    let (c, e, f, h, m, s) = get_time_tuple(ts, tz_offset);
471    format!("{:>04}-{:>02}-{:>02} {:>02}:{:>02}:{:>02}", c, e, f, h, m, s)
472}
473
474pub fn to_human_readable_string(size: u64) -> String {
475    if size < 1024 {
476        return format!("{}", size);
477    }
478
479    static UNITS: [&'static str; 7] = ["", "K", "M", "G", "T", "P", "E"];
480
481    let sizef = size as f64;
482    let digit_groups = (sizef.log10() / 1024f64.log10()) as i32;
483    format!("{:.1}{}",
484            sizef / 1024f64.powf(digit_groups as f64),
485            UNITS[digit_groups as usize])
486}
487
488#[cfg(test)]
489mod tests {
490    use super::ArgParser;
491
492    #[test]
493    fn stop_parsing() {
494        let args = vec![String::from("binname"), String::from("-a"), String::from("--"), String::from("-v")];
495        let mut parser = ArgParser::new(2);
496        parser = parser.add_flag(&["a"]).add_flag(&["v"]);
497        parser.parse(args.into_iter());
498        assert!(parser.found(&'a'));
499        assert!(!parser.found(&'v'));
500        assert!(parser.args[0] == "-v");
501    }
502
503    #[test]
504    fn short_opts() {
505        let args = vec![String::from("binname"), String::from("-asdf"), String::from("-f"), String::from("foo")];
506        let mut parser = ArgParser::new(4);
507        parser = parser.add_flag(&["a"])
508            .add_flag(&["d"])
509            .add_opt("s", "")
510            .add_opt_default("w", "", "default")
511            .add_opt("f", "");
512        parser.parse(args.into_iter());
513        assert!(parser.found(&'a'));
514        assert!(!parser.found(&'d'));
515        assert!(parser.get_opt(&'s') == Some(String::from("df")));
516        assert!(parser.get_opt(&'w') == Some(String::from("default")));
517        assert!(parser.get_opt(&'f') == Some(String::from("foo")));
518    }
519
520    #[test]
521    fn long_opts() {
522        let args = vec![String::from("binname"), String::from("--foo=bar")];
523        let mut parser = ArgParser::new(4);
524        parser = parser.add_opt("", "foo");
525        parser = parser.add_opt_default("", "def", "default");
526        parser.parse(args.into_iter());
527        assert!(parser.get_opt("foo") == Some(String::from("bar")));
528        assert!(parser.get_opt("def") == Some(String::from("default")));
529    }
530
531    #[test]
532    fn settings() {
533        let args = vec![String::from("binname"), String::from("-h"), String::from("if=bar")];
534        let mut parser = ArgParser::new(4);
535        parser = parser.add_flag(&["h"]).add_setting("if").add_setting_default("of", "foo");
536        parser.parse(args.into_iter());
537        assert!(parser.found("if"));
538        assert!(parser.get_setting("if") == Some(String::from("bar")));
539        assert!(parser.get_setting("of") == Some(String::from("foo")));
540    }
541}