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)]
9pub 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)]
47struct Rhs<T> {
49 value: T,
51 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)]
65enum Value {
67 Flag(Rhs<Rc<RefCell<bool>>>),
68 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#[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 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 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 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 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 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 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 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 let arg = &arg[2..];
228 if arg.is_empty() {
229 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..]; 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 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..]; 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 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 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 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 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 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 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; 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
444pub 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}