1use std::collections::BTreeMap;
13
14#[derive(Debug, Clone, Default)]
19pub struct ArgsSpec {
20 pub options: Vec<OptionDecl>,
21 pub positionals: Vec<PositionalDecl>,
22 pub lenient_unknowns: bool,
32}
33
34#[derive(Debug, Clone)]
36pub struct OptionDecl {
37 pub long: String,
39 pub short: Option<char>,
41 pub takes_value: bool,
43 pub allows_bare: bool,
46 pub repeatable: bool,
49 pub choices: Option<Vec<String>>,
51}
52
53#[derive(Debug, Clone)]
55pub struct PositionalDecl {
56 pub name: String,
58 pub required: bool,
60 pub variadic: bool,
62 pub choices: Option<Vec<String>>,
64}
65
66#[derive(Debug, Default)]
71pub struct ParsedArgs {
72 pub options: BTreeMap<String, OptionState>,
74 pub positionals: Vec<String>,
76}
77
78#[derive(Debug, Clone)]
80pub enum OptionState {
81 BarePresent,
83 WithValues(Vec<String>),
85}
86
87pub fn parse(spec: &ArgsSpec, args: &[String]) -> Result<ParsedArgs, String> {
94 let mut out = ParsedArgs::default();
95 let mut i = 1usize;
96 let mut stop_flags = false;
97
98 while i < args.len() {
99 let tok = &args[i];
100
101 if !stop_flags && tok == "--" {
102 stop_flags = true;
103 i += 1;
104 continue;
105 }
106
107 if !stop_flags && tok.starts_with("--") {
108 let rest = &tok[2..];
110 let (name, inline_value) = match rest.split_once('=') {
111 Some((n, v)) => (n, Some(v.to_string())),
112 None => (rest, None),
113 };
114 match find_long(spec, name) {
115 Some(decl) => {
116 i = consume_flag(decl, inline_value, args, i, &mut out)?;
117 }
118 None if spec.lenient_unknowns => {
119 i += 1;
125 }
126 None => return Err(unknown_long_error(spec, name)),
127 }
128 continue;
129 }
130
131 if !stop_flags && tok.starts_with('-') && tok.len() >= 2 {
132 let rest = &tok[1..];
134 if let Some((head, inline_value)) = rest.split_once('=') {
135 if head.chars().count() != 1 {
137 return Err(format!(
138 "invalid short-flag syntax `{tok}`: `-x=value` requires a single-letter short"
139 ));
140 }
141 let c = head.chars().next().unwrap();
142 match find_short(spec, c) {
143 Some(decl) => {
144 i = consume_flag(decl, Some(inline_value.to_string()), args, i, &mut out)?;
145 }
146 None if spec.lenient_unknowns => {
147 i += 1;
148 }
149 None => return Err(format!("unknown short flag `-{c}`")),
150 }
151 continue;
152 }
153 let chars: Vec<char> = rest.chars().collect();
157 if spec.lenient_unknowns && chars.iter().any(|c| find_short(spec, *c).is_none()) {
158 i += 1;
163 continue;
164 }
165 for (pos, c) in chars.iter().enumerate() {
166 let decl =
167 find_short(spec, *c).ok_or_else(|| format!("unknown short flag `-{c}`"))?;
168 let is_last = pos == chars.len() - 1;
169 if !is_last && decl.takes_value {
170 return Err(format!(
171 "short `-{c}` takes a value and cannot be clustered mid-token `{tok}`"
172 ));
173 }
174 if is_last {
175 i = consume_flag(decl, None, args, i, &mut out)?;
176 } else {
177 record_option(&mut out, decl, None, spec)?;
178 }
179 }
180 if chars.is_empty() {
181 out.positionals.push(tok.clone());
183 i += 1;
184 }
185 continue;
186 }
187
188 out.positionals.push(tok.clone());
190 i += 1;
191 }
192
193 if !spec.lenient_unknowns {
198 let required_count = spec.positionals.iter().filter(|p| p.required).count();
199 if out.positionals.len() < required_count {
200 let missing = &spec.positionals[out.positionals.len()].name;
201 return Err(format!("missing required argument <{missing}>"));
202 }
203 }
204
205 for (idx, value) in out.positionals.iter().enumerate() {
213 match positional_decl_for(spec, idx) {
214 Some(d) => {
215 if let Some(choices) = &d.choices
216 && !choices.iter().any(|c| c == value)
217 {
218 return Err(format!(
219 "invalid value `{value}` for <{}> -- allowed: {}",
220 d.name,
221 choices.join(", ")
222 ));
223 }
224 }
225 None if spec.lenient_unknowns => {}
226 None => {
227 let hint = if value.starts_with('-') {
228 "; tokens after a standalone `--` are treated as \
229 positional arguments, not options — pass options \
230 directly, without a `--` separator"
231 } else {
232 ""
233 };
234 return Err(format!("unexpected argument `{value}`{hint}"));
235 }
236 }
237 }
238
239 Ok(out)
240}
241
242fn is_flag_like(s: &str) -> bool {
253 s.starts_with('-') && s != "-" && s != "--" && s.parse::<f64>().is_err()
254}
255
256fn consume_flag(
257 decl: &OptionDecl,
258 inline_value: Option<String>,
259 args: &[String],
260 i: usize,
261 out: &mut ParsedArgs,
262) -> Result<usize, String> {
263 if !decl.takes_value {
264 if inline_value.is_some() {
266 return Err(format!("flag `--{}` takes no value", decl.long));
267 }
268 record_option(out, decl, None, &ArgsSpec::default())?;
269 return Ok(i + 1);
270 }
271
272 if let Some(v) = inline_value {
274 record_option(out, decl, Some(v), &ArgsSpec::default())?;
275 return Ok(i + 1);
276 }
277
278 let next_idx = i + 1;
280 let next_is_flag = args.get(next_idx).map(|s| is_flag_like(s)).unwrap_or(true); if !next_is_flag {
283 let v = args[next_idx].clone();
284 record_option(out, decl, Some(v), &ArgsSpec::default())?;
285 return Ok(i + 2);
286 }
287
288 if !decl.allows_bare {
290 return Err(format!("`--{}` requires a value", decl.long));
291 }
292 record_option(out, decl, None, &ArgsSpec::default())?;
293 Ok(i + 1)
294}
295
296fn record_option(
299 out: &mut ParsedArgs,
300 decl: &OptionDecl,
301 value: Option<String>,
302 _spec: &ArgsSpec,
303) -> Result<(), String> {
304 if let (Some(v), Some(choices)) = (&value, &decl.choices) {
306 for part in split_list_value(v) {
307 if !choices.iter().any(|c| c == part) {
308 return Err(format!(
309 "invalid value `{part}` for `--{}` -- allowed: {}",
310 decl.long,
311 choices.join(", ")
312 ));
313 }
314 }
315 }
316
317 let key = decl.long.clone();
318 match (value, decl.repeatable) {
319 (None, _) => {
320 out.options.entry(key).or_insert(OptionState::BarePresent);
322 }
323 (Some(v), false) => {
324 out.options.insert(key, OptionState::WithValues(vec![v]));
326 }
327 (Some(v), true) => {
328 let parts: Vec<String> = split_list_value(&v).into_iter().map(String::from).collect();
330 let entry = out
331 .options
332 .entry(key)
333 .or_insert(OptionState::WithValues(Vec::new()));
334 if let OptionState::WithValues(list) = entry {
335 list.extend(parts);
336 }
337 }
338 }
339 Ok(())
340}
341
342fn split_list_value(v: &str) -> Vec<&str> {
345 v.split(',')
346 .map(str::trim)
347 .filter(|s| !s.is_empty())
348 .collect()
349}
350
351fn find_long<'a>(spec: &'a ArgsSpec, name: &str) -> Option<&'a OptionDecl> {
352 spec.options.iter().find(|o| o.long == name)
353}
354
355fn find_short(spec: &ArgsSpec, c: char) -> Option<&OptionDecl> {
356 spec.options.iter().find(|o| o.short == Some(c))
357}
358
359fn positional_decl_for(spec: &ArgsSpec, idx: usize) -> Option<&PositionalDecl> {
360 if let Some(decl) = spec.positionals.get(idx) {
362 return Some(decl);
363 }
364 spec.positionals.iter().rev().find(|d| d.variadic)
365}
366
367fn unknown_long_error(spec: &ArgsSpec, name: &str) -> String {
369 let suggestion = spec
370 .options
371 .iter()
372 .filter(|o| similar(&o.long, name))
373 .map(|o| format!("--{}", o.long))
374 .next();
375 match suggestion {
376 Some(s) => format!("unknown flag `--{name}`, did you mean `{s}`?"),
377 None => format!("unknown flag `--{name}`"),
378 }
379}
380
381pub fn suggest(candidates: &[&str], input: &str) -> Option<String> {
387 candidates
388 .iter()
389 .find(|c| similar(c, input))
390 .map(|c| (*c).to_string())
391}
392
393fn similar(candidate: &str, target: &str) -> bool {
398 if candidate == target {
399 return false;
400 }
401 levenshtein(candidate, target) <= 2
402}
403
404fn levenshtein(a: &str, b: &str) -> usize {
405 let a: Vec<char> = a.chars().collect();
406 let b: Vec<char> = b.chars().collect();
407 let (m, n) = (a.len(), b.len());
408 if m == 0 {
409 return n;
410 }
411 if n == 0 {
412 return m;
413 }
414 let mut prev: Vec<usize> = (0..=n).collect();
415 let mut curr = vec![0usize; n + 1];
416 for (i, ca) in a.iter().enumerate() {
417 curr[0] = i + 1;
418 for (j, cb) in b.iter().enumerate() {
419 let cost = if ca == cb { 0 } else { 1 };
420 curr[j + 1] = (prev[j + 1] + 1).min(curr[j] + 1).min(prev[j] + cost);
421 }
422 std::mem::swap(&mut prev, &mut curr);
423 }
424 prev[n]
425}
426
427#[cfg(test)]
430mod tests {
431 use super::*;
432
433 fn flag(long: &str, short: Option<char>) -> OptionDecl {
434 OptionDecl {
435 long: long.into(),
436 short,
437 takes_value: false,
438 allows_bare: true,
439 repeatable: false,
440 choices: None,
441 }
442 }
443
444 fn value(long: &str, short: Option<char>, bare_ok: bool) -> OptionDecl {
445 OptionDecl {
446 long: long.into(),
447 short,
448 takes_value: true,
449 allows_bare: bare_ok,
450 repeatable: false,
451 choices: None,
452 }
453 }
454
455 fn list(long: &str, short: Option<char>) -> OptionDecl {
456 OptionDecl {
457 long: long.into(),
458 short,
459 takes_value: true,
460 allows_bare: false,
461 repeatable: true,
462 choices: None,
463 }
464 }
465
466 fn pos(name: &str, required: bool, variadic: bool) -> PositionalDecl {
467 PositionalDecl {
468 name: name.into(),
469 required,
470 variadic,
471 choices: None,
472 }
473 }
474
475 fn argv(parts: &[&str]) -> Vec<String> {
476 std::iter::once("prog")
477 .chain(parts.iter().copied())
478 .map(String::from)
479 .collect()
480 }
481
482 #[test]
483 fn parses_long_flag_with_value() {
484 let spec = ArgsSpec {
485 options: vec![value("model", None, false)],
486 positionals: vec![],
487 ..ArgsSpec::default()
488 };
489 let out = parse(&spec, &argv(&["--model", "mlp"])).unwrap();
490 match out.options.get("model") {
491 Some(OptionState::WithValues(v)) => assert_eq!(v, &vec!["mlp".to_string()]),
492 other => panic!("expected WithValues, got {:?}", other),
493 }
494 }
495
496 #[test]
497 fn is_flag_like_treats_negative_numbers_as_values() {
498 assert!(is_flag_like("--verbose"));
501 assert!(is_flag_like("-v"));
502 assert!(!is_flag_like("-0.5"));
503 assert!(!is_flag_like("-5"));
504 assert!(!is_flag_like("-1e-3"));
505 assert!(!is_flag_like("-")); assert!(!is_flag_like("--")); assert!(!is_flag_like("mlp")); }
509
510 #[test]
511 fn value_flag_consumes_space_separated_negative_number() {
512 let spec = ArgsSpec {
514 options: vec![value("lr", None, false)],
515 positionals: vec![],
516 ..ArgsSpec::default()
517 };
518 let out = parse(&spec, &argv(&["--lr", "-0.5"])).unwrap();
519 match out.options.get("lr") {
520 Some(OptionState::WithValues(v)) => assert_eq!(v, &vec!["-0.5".to_string()]),
521 other => panic!("expected WithValues([-0.5]), got {:?}", other),
522 }
523 }
524
525 #[test]
526 fn value_flag_still_errors_when_next_is_a_real_flag() {
527 let spec = ArgsSpec {
529 options: vec![value("lr", None, false), flag("verbose", None)],
530 positionals: vec![],
531 ..ArgsSpec::default()
532 };
533 let err = parse(&spec, &argv(&["--lr", "--verbose"])).unwrap_err();
534 assert!(err.contains("requires a value"), "got: {err}");
535 }
536
537 #[test]
538 fn parses_long_flag_with_equals() {
539 let spec = ArgsSpec {
540 options: vec![value("model", None, false)],
541 positionals: vec![],
542 ..ArgsSpec::default()
543 };
544 let out = parse(&spec, &argv(&["--model=mlp"])).unwrap();
545 match out.options.get("model") {
546 Some(OptionState::WithValues(v)) => assert_eq!(v, &vec!["mlp".to_string()]),
547 _ => panic!("expected WithValues"),
548 }
549 }
550
551 #[test]
552 fn bare_flag_without_default_errors() {
553 let spec = ArgsSpec {
554 options: vec![value("report", None, false)],
555 positionals: vec![],
556 ..ArgsSpec::default()
557 };
558 let err = parse(&spec, &argv(&["--report"])).unwrap_err();
559 assert!(err.contains("requires a value"), "got: {err}");
560 }
561
562 #[test]
563 fn bare_flag_with_default_is_present() {
564 let spec = ArgsSpec {
565 options: vec![value("report", None, true)],
566 positionals: vec![],
567 ..ArgsSpec::default()
568 };
569 let out = parse(&spec, &argv(&["--report"])).unwrap();
570 assert!(matches!(
571 out.options.get("report"),
572 Some(OptionState::BarePresent)
573 ));
574 }
575
576 #[test]
577 fn bool_flag_presence() {
578 let spec = ArgsSpec {
579 options: vec![flag("validate", None)],
580 positionals: vec![],
581 ..ArgsSpec::default()
582 };
583 let out = parse(&spec, &argv(&["--validate"])).unwrap();
584 assert!(matches!(
585 out.options.get("validate"),
586 Some(OptionState::BarePresent)
587 ));
588 }
589
590 #[test]
591 fn bool_flag_rejects_value() {
592 let spec = ArgsSpec {
593 options: vec![flag("validate", None)],
594 positionals: vec![],
595 ..ArgsSpec::default()
596 };
597 let err = parse(&spec, &argv(&["--validate=yes"])).unwrap_err();
598 assert!(err.contains("takes no value"), "got: {err}");
599 }
600
601 #[test]
602 fn short_flag() {
603 let spec = ArgsSpec {
604 options: vec![flag("verbose", Some('v'))],
605 positionals: vec![],
606 ..ArgsSpec::default()
607 };
608 let out = parse(&spec, &argv(&["-v"])).unwrap();
609 assert!(matches!(
610 out.options.get("verbose"),
611 Some(OptionState::BarePresent)
612 ));
613 }
614
615 #[test]
616 fn short_clustering_for_bool_flags() {
617 let spec = ArgsSpec {
618 options: vec![flag("a", Some('a')), flag("b", Some('b'))],
619 positionals: vec![],
620 ..ArgsSpec::default()
621 };
622 let out = parse(&spec, &argv(&["-ab"])).unwrap();
623 assert!(out.options.contains_key("a"));
624 assert!(out.options.contains_key("b"));
625 }
626
627 #[test]
628 fn short_cluster_last_may_take_value() {
629 let spec = ArgsSpec {
630 options: vec![flag("a", Some('a')), value("model", Some('m'), false)],
631 positionals: vec![],
632 ..ArgsSpec::default()
633 };
634 let out = parse(&spec, &argv(&["-am", "mlp"])).unwrap();
635 assert!(out.options.contains_key("a"));
636 match out.options.get("model") {
637 Some(OptionState::WithValues(v)) => assert_eq!(v, &vec!["mlp".to_string()]),
638 _ => panic!("expected model value"),
639 }
640 }
641
642 #[test]
643 fn list_option_accumulates_across_repeats_and_commas() {
644 let spec = ArgsSpec {
645 options: vec![list("tags", Some('t'))],
646 positionals: vec![],
647 ..ArgsSpec::default()
648 };
649 let out = parse(&spec, &argv(&["--tags", "a,b", "-t", "c"])).unwrap();
650 match out.options.get("tags") {
651 Some(OptionState::WithValues(v)) => {
652 assert_eq!(v, &vec!["a".to_string(), "b".into(), "c".into()]);
653 }
654 _ => panic!("expected list values"),
655 }
656 }
657
658 #[test]
659 fn positionals_in_order() {
660 let spec = ArgsSpec {
661 options: vec![],
662 positionals: vec![pos("first", true, false), pos("second", false, false)],
663 ..ArgsSpec::default()
664 };
665 let out = parse(&spec, &argv(&["a", "b"])).unwrap();
666 assert_eq!(out.positionals, vec!["a".to_string(), "b".into()]);
667 }
668
669 #[test]
670 fn missing_required_positional_errors() {
671 let spec = ArgsSpec {
672 options: vec![],
673 positionals: vec![pos("first", true, false)],
674 ..ArgsSpec::default()
675 };
676 let err = parse(&spec, &argv(&[])).unwrap_err();
677 assert!(err.contains("missing required argument"), "got: {err}");
678 }
679
680 #[test]
681 fn variadic_positional_absorbs_tail() {
682 let spec = ArgsSpec {
683 options: vec![],
684 positionals: vec![pos("files", false, true)],
685 ..ArgsSpec::default()
686 };
687 let out = parse(&spec, &argv(&["a", "b", "c"])).unwrap();
688 assert_eq!(
689 out.positionals,
690 vec!["a".to_string(), "b".into(), "c".into()]
691 );
692 }
693
694 #[test]
695 fn double_dash_stops_flag_parsing() {
696 let spec = ArgsSpec {
697 options: vec![flag("verbose", None)],
698 positionals: vec![pos("rest", false, true)],
699 ..ArgsSpec::default()
700 };
701 let out = parse(&spec, &argv(&["--", "--verbose", "-x"])).unwrap();
702 assert!(!out.options.contains_key("verbose"));
703 assert_eq!(out.positionals, vec!["--verbose".to_string(), "-x".into()]);
704 }
705
706 #[test]
707 fn excess_positional_errors_loudly() {
708 let spec = ArgsSpec {
710 options: vec![value("model", None, false)],
711 positionals: vec![],
712 ..ArgsSpec::default()
713 };
714 let err = parse(&spec, &argv(&["stray"])).unwrap_err();
715 assert!(err.contains("unexpected argument `stray`"), "got: {err}");
716 }
717
718 #[test]
719 fn dashdash_forwarded_options_error_with_hint() {
720 let spec = ArgsSpec {
724 options: vec![value("model", None, false)],
725 positionals: vec![],
726 ..ArgsSpec::default()
727 };
728 let err = parse(&spec, &argv(&["--", "--model", "lenet"])).unwrap_err();
729 assert!(err.contains("unexpected argument `--model`"), "got: {err}");
730 assert!(err.contains("without a `--` separator"), "got: {err}");
731 }
732
733 #[test]
734 fn lenient_mode_tolerates_excess_positionals() {
735 let spec = ArgsSpec {
739 options: vec![],
740 positionals: vec![],
741 lenient_unknowns: true,
742 };
743 let out = parse(&spec, &argv(&["--unknown", "orphan-value"])).unwrap();
744 assert_eq!(out.positionals, vec!["orphan-value".to_string()]);
745 }
746
747 #[test]
748 fn unknown_flag_suggests_similar() {
749 let spec = ArgsSpec {
750 options: vec![value("model", None, false)],
751 positionals: vec![],
752 ..ArgsSpec::default()
753 };
754 let err = parse(&spec, &argv(&["--modl", "mlp"])).unwrap_err();
755 assert!(err.contains("did you mean"), "got: {err}");
756 }
757
758 #[test]
759 fn choices_validated_at_parse_time() {
760 let mut model = value("model", None, false);
761 model.choices = Some(vec!["mlp".into(), "lenet".into()]);
762 let spec = ArgsSpec {
763 options: vec![model],
764 positionals: vec![],
765 ..ArgsSpec::default()
766 };
767 let err = parse(&spec, &argv(&["--model", "foobar"])).unwrap_err();
768 assert!(err.contains("allowed"), "got: {err}");
769 }
770
771 #[test]
772 fn bare_dash_is_positional() {
773 let spec = ArgsSpec {
774 options: vec![],
775 positionals: vec![pos("target", true, false)],
776 ..ArgsSpec::default()
777 };
778 let out = parse(&spec, &argv(&["-"])).unwrap();
779 assert_eq!(out.positionals, vec!["-".to_string()]);
780 }
781
782 #[test]
783 fn scalar_last_write_wins() {
784 let spec = ArgsSpec {
785 options: vec![value("model", None, false)],
786 positionals: vec![],
787 ..ArgsSpec::default()
788 };
789 let out = parse(&spec, &argv(&["--model", "a", "--model", "b"])).unwrap();
790 match out.options.get("model") {
791 Some(OptionState::WithValues(v)) => assert_eq!(v, &vec!["b".to_string()]),
792 _ => panic!("expected last-write-wins"),
793 }
794 }
795}