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 = find_short(spec, *c)
167 .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 if !choices.iter().any(|c| c == value) {
217 return Err(format!(
218 "invalid value `{value}` for <{}> -- allowed: {}",
219 d.name,
220 choices.join(", ")
221 ));
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
281 .get(next_idx)
282 .map(|s| is_flag_like(s))
283 .unwrap_or(true); if !next_is_flag {
286 let v = args[next_idx].clone();
287 record_option(out, decl, Some(v), &ArgsSpec::default())?;
288 return Ok(i + 2);
289 }
290
291 if !decl.allows_bare {
293 return Err(format!("`--{}` requires a value", decl.long));
294 }
295 record_option(out, decl, None, &ArgsSpec::default())?;
296 Ok(i + 1)
297}
298
299fn record_option(
302 out: &mut ParsedArgs,
303 decl: &OptionDecl,
304 value: Option<String>,
305 _spec: &ArgsSpec,
306) -> Result<(), String> {
307 if let (Some(v), Some(choices)) = (&value, &decl.choices) {
309 for part in split_list_value(v) {
310 if !choices.iter().any(|c| c == part) {
311 return Err(format!(
312 "invalid value `{part}` for `--{}` -- allowed: {}",
313 decl.long,
314 choices.join(", ")
315 ));
316 }
317 }
318 }
319
320 let key = decl.long.clone();
321 match (value, decl.repeatable) {
322 (None, _) => {
323 out.options.entry(key).or_insert(OptionState::BarePresent);
325 }
326 (Some(v), false) => {
327 out.options.insert(key, OptionState::WithValues(vec![v]));
329 }
330 (Some(v), true) => {
331 let parts: Vec<String> = split_list_value(&v).into_iter().map(String::from).collect();
333 let entry = out
334 .options
335 .entry(key)
336 .or_insert(OptionState::WithValues(Vec::new()));
337 if let OptionState::WithValues(list) = entry {
338 list.extend(parts);
339 }
340 }
341 }
342 Ok(())
343}
344
345fn split_list_value(v: &str) -> Vec<&str> {
348 v.split(',').map(str::trim).filter(|s| !s.is_empty()).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)
421 .min(curr[j] + 1)
422 .min(prev[j] + cost);
423 }
424 std::mem::swap(&mut prev, &mut curr);
425 }
426 prev[n]
427}
428
429#[cfg(test)]
432mod tests {
433 use super::*;
434
435 fn flag(long: &str, short: Option<char>) -> OptionDecl {
436 OptionDecl {
437 long: long.into(),
438 short,
439 takes_value: false,
440 allows_bare: true,
441 repeatable: false,
442 choices: None,
443 }
444 }
445
446 fn value(long: &str, short: Option<char>, bare_ok: bool) -> OptionDecl {
447 OptionDecl {
448 long: long.into(),
449 short,
450 takes_value: true,
451 allows_bare: bare_ok,
452 repeatable: false,
453 choices: None,
454 }
455 }
456
457 fn list(long: &str, short: Option<char>) -> OptionDecl {
458 OptionDecl {
459 long: long.into(),
460 short,
461 takes_value: true,
462 allows_bare: false,
463 repeatable: true,
464 choices: None,
465 }
466 }
467
468 fn pos(name: &str, required: bool, variadic: bool) -> PositionalDecl {
469 PositionalDecl {
470 name: name.into(),
471 required,
472 variadic,
473 choices: None,
474 }
475 }
476
477 fn argv(parts: &[&str]) -> Vec<String> {
478 std::iter::once("prog")
479 .chain(parts.iter().copied())
480 .map(String::from)
481 .collect()
482 }
483
484 #[test]
485 fn parses_long_flag_with_value() {
486 let spec = ArgsSpec {
487 options: vec![value("model", None, false)],
488 positionals: vec![],
489 ..ArgsSpec::default()
490 };
491 let out = parse(&spec, &argv(&["--model", "mlp"])).unwrap();
492 match out.options.get("model") {
493 Some(OptionState::WithValues(v)) => assert_eq!(v, &vec!["mlp".to_string()]),
494 other => panic!("expected WithValues, got {:?}", other),
495 }
496 }
497
498 #[test]
499 fn is_flag_like_treats_negative_numbers_as_values() {
500 assert!(is_flag_like("--verbose"));
503 assert!(is_flag_like("-v"));
504 assert!(!is_flag_like("-0.5"));
505 assert!(!is_flag_like("-5"));
506 assert!(!is_flag_like("-1e-3"));
507 assert!(!is_flag_like("-")); assert!(!is_flag_like("--")); assert!(!is_flag_like("mlp")); }
511
512 #[test]
513 fn value_flag_consumes_space_separated_negative_number() {
514 let spec = ArgsSpec {
516 options: vec![value("lr", None, false)],
517 positionals: vec![],
518 ..ArgsSpec::default()
519 };
520 let out = parse(&spec, &argv(&["--lr", "-0.5"])).unwrap();
521 match out.options.get("lr") {
522 Some(OptionState::WithValues(v)) => assert_eq!(v, &vec!["-0.5".to_string()]),
523 other => panic!("expected WithValues([-0.5]), got {:?}", other),
524 }
525 }
526
527 #[test]
528 fn value_flag_still_errors_when_next_is_a_real_flag() {
529 let spec = ArgsSpec {
531 options: vec![value("lr", None, false), flag("verbose", None)],
532 positionals: vec![],
533 ..ArgsSpec::default()
534 };
535 let err = parse(&spec, &argv(&["--lr", "--verbose"])).unwrap_err();
536 assert!(err.contains("requires a value"), "got: {err}");
537 }
538
539 #[test]
540 fn parses_long_flag_with_equals() {
541 let spec = ArgsSpec {
542 options: vec![value("model", None, false)],
543 positionals: vec![],
544 ..ArgsSpec::default()
545 };
546 let out = parse(&spec, &argv(&["--model=mlp"])).unwrap();
547 match out.options.get("model") {
548 Some(OptionState::WithValues(v)) => assert_eq!(v, &vec!["mlp".to_string()]),
549 _ => panic!("expected WithValues"),
550 }
551 }
552
553 #[test]
554 fn bare_flag_without_default_errors() {
555 let spec = ArgsSpec {
556 options: vec![value("report", None, false)],
557 positionals: vec![],
558 ..ArgsSpec::default()
559 };
560 let err = parse(&spec, &argv(&["--report"])).unwrap_err();
561 assert!(err.contains("requires a value"), "got: {err}");
562 }
563
564 #[test]
565 fn bare_flag_with_default_is_present() {
566 let spec = ArgsSpec {
567 options: vec![value("report", None, true)],
568 positionals: vec![],
569 ..ArgsSpec::default()
570 };
571 let out = parse(&spec, &argv(&["--report"])).unwrap();
572 assert!(matches!(out.options.get("report"), Some(OptionState::BarePresent)));
573 }
574
575 #[test]
576 fn bool_flag_presence() {
577 let spec = ArgsSpec {
578 options: vec![flag("validate", None)],
579 positionals: vec![],
580 ..ArgsSpec::default()
581 };
582 let out = parse(&spec, &argv(&["--validate"])).unwrap();
583 assert!(matches!(out.options.get("validate"), Some(OptionState::BarePresent)));
584 }
585
586 #[test]
587 fn bool_flag_rejects_value() {
588 let spec = ArgsSpec {
589 options: vec![flag("validate", None)],
590 positionals: vec![],
591 ..ArgsSpec::default()
592 };
593 let err = parse(&spec, &argv(&["--validate=yes"])).unwrap_err();
594 assert!(err.contains("takes no value"), "got: {err}");
595 }
596
597 #[test]
598 fn short_flag() {
599 let spec = ArgsSpec {
600 options: vec![flag("verbose", Some('v'))],
601 positionals: vec![],
602 ..ArgsSpec::default()
603 };
604 let out = parse(&spec, &argv(&["-v"])).unwrap();
605 assert!(matches!(out.options.get("verbose"), Some(OptionState::BarePresent)));
606 }
607
608 #[test]
609 fn short_clustering_for_bool_flags() {
610 let spec = ArgsSpec {
611 options: vec![flag("a", Some('a')), flag("b", Some('b'))],
612 positionals: vec![],
613 ..ArgsSpec::default()
614 };
615 let out = parse(&spec, &argv(&["-ab"])).unwrap();
616 assert!(out.options.contains_key("a"));
617 assert!(out.options.contains_key("b"));
618 }
619
620 #[test]
621 fn short_cluster_last_may_take_value() {
622 let spec = ArgsSpec {
623 options: vec![flag("a", Some('a')), value("model", Some('m'), false)],
624 positionals: vec![],
625 ..ArgsSpec::default()
626 };
627 let out = parse(&spec, &argv(&["-am", "mlp"])).unwrap();
628 assert!(out.options.contains_key("a"));
629 match out.options.get("model") {
630 Some(OptionState::WithValues(v)) => assert_eq!(v, &vec!["mlp".to_string()]),
631 _ => panic!("expected model value"),
632 }
633 }
634
635 #[test]
636 fn list_option_accumulates_across_repeats_and_commas() {
637 let spec = ArgsSpec {
638 options: vec![list("tags", Some('t'))],
639 positionals: vec![],
640 ..ArgsSpec::default()
641 };
642 let out = parse(&spec, &argv(&["--tags", "a,b", "-t", "c"])).unwrap();
643 match out.options.get("tags") {
644 Some(OptionState::WithValues(v)) => {
645 assert_eq!(v, &vec!["a".to_string(), "b".into(), "c".into()]);
646 }
647 _ => panic!("expected list values"),
648 }
649 }
650
651 #[test]
652 fn positionals_in_order() {
653 let spec = ArgsSpec {
654 options: vec![],
655 positionals: vec![pos("first", true, false), pos("second", false, false)],
656 ..ArgsSpec::default()
657 };
658 let out = parse(&spec, &argv(&["a", "b"])).unwrap();
659 assert_eq!(out.positionals, vec!["a".to_string(), "b".into()]);
660 }
661
662 #[test]
663 fn missing_required_positional_errors() {
664 let spec = ArgsSpec {
665 options: vec![],
666 positionals: vec![pos("first", true, false)],
667 ..ArgsSpec::default()
668 };
669 let err = parse(&spec, &argv(&[])).unwrap_err();
670 assert!(err.contains("missing required argument"), "got: {err}");
671 }
672
673 #[test]
674 fn variadic_positional_absorbs_tail() {
675 let spec = ArgsSpec {
676 options: vec![],
677 positionals: vec![pos("files", false, true)],
678 ..ArgsSpec::default()
679 };
680 let out = parse(&spec, &argv(&["a", "b", "c"])).unwrap();
681 assert_eq!(out.positionals, vec!["a".to_string(), "b".into(), "c".into()]);
682 }
683
684 #[test]
685 fn double_dash_stops_flag_parsing() {
686 let spec = ArgsSpec {
687 options: vec![flag("verbose", None)],
688 positionals: vec![pos("rest", false, true)],
689 ..ArgsSpec::default()
690 };
691 let out = parse(&spec, &argv(&["--", "--verbose", "-x"])).unwrap();
692 assert!(!out.options.contains_key("verbose"));
693 assert_eq!(out.positionals, vec!["--verbose".to_string(), "-x".into()]);
694 }
695
696 #[test]
697 fn excess_positional_errors_loudly() {
698 let spec = ArgsSpec {
700 options: vec![value("model", None, false)],
701 positionals: vec![],
702 ..ArgsSpec::default()
703 };
704 let err = parse(&spec, &argv(&["stray"])).unwrap_err();
705 assert!(err.contains("unexpected argument `stray`"), "got: {err}");
706 }
707
708 #[test]
709 fn dashdash_forwarded_options_error_with_hint() {
710 let spec = ArgsSpec {
714 options: vec![value("model", None, false)],
715 positionals: vec![],
716 ..ArgsSpec::default()
717 };
718 let err = parse(&spec, &argv(&["--", "--model", "lenet"])).unwrap_err();
719 assert!(err.contains("unexpected argument `--model`"), "got: {err}");
720 assert!(err.contains("without a `--` separator"), "got: {err}");
721 }
722
723 #[test]
724 fn lenient_mode_tolerates_excess_positionals() {
725 let spec = ArgsSpec {
729 options: vec![],
730 positionals: vec![],
731 lenient_unknowns: true,
732 };
733 let out = parse(&spec, &argv(&["--unknown", "orphan-value"])).unwrap();
734 assert_eq!(out.positionals, vec!["orphan-value".to_string()]);
735 }
736
737 #[test]
738 fn unknown_flag_suggests_similar() {
739 let spec = ArgsSpec {
740 options: vec![value("model", None, false)],
741 positionals: vec![],
742 ..ArgsSpec::default()
743 };
744 let err = parse(&spec, &argv(&["--modl", "mlp"])).unwrap_err();
745 assert!(err.contains("did you mean"), "got: {err}");
746 }
747
748 #[test]
749 fn choices_validated_at_parse_time() {
750 let mut model = value("model", None, false);
751 model.choices = Some(vec!["mlp".into(), "lenet".into()]);
752 let spec = ArgsSpec {
753 options: vec![model],
754 positionals: vec![],
755 ..ArgsSpec::default()
756 };
757 let err = parse(&spec, &argv(&["--model", "foobar"])).unwrap_err();
758 assert!(err.contains("allowed"), "got: {err}");
759 }
760
761 #[test]
762 fn bare_dash_is_positional() {
763 let spec = ArgsSpec {
764 options: vec![],
765 positionals: vec![pos("target", true, false)],
766 ..ArgsSpec::default()
767 };
768 let out = parse(&spec, &argv(&["-"])).unwrap();
769 assert_eq!(out.positionals, vec!["-".to_string()]);
770 }
771
772 #[test]
773 fn scalar_last_write_wins() {
774 let spec = ArgsSpec {
775 options: vec![value("model", None, false)],
776 positionals: vec![],
777 ..ArgsSpec::default()
778 };
779 let out = parse(&spec, &argv(&["--model", "a", "--model", "b"])).unwrap();
780 match out.options.get("model") {
781 Some(OptionState::WithValues(v)) => assert_eq!(v, &vec!["b".to_string()]),
782 _ => panic!("expected last-write-wins"),
783 }
784 }
785}