missing string

1 files changed, 26 insertions, 3 deletions

diff --git a/src/lib.rs b/src/lib.rs
index 9f533b9..de54b89 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,4 +1,4 @@
-use std::str::FromStr;
+use std::{borrow::Cow, cell::Cell, str::FromStr};
 
 use formula::Formula;
 
@@ -10,6 +10,7 @@ pub struct Scurvy {
 	unknown: Vec<(String, String)>,
 	print_help: bool,
 	print_version: bool,
+	missing_string: Cell<Cow<'static, str>>,
 }
 
 struct Pair {
@@ -51,6 +52,7 @@ impl Scurvy {
 			unknown,
 			print_help,
 			print_version,
+			missing_string: Cell::new(Cow::Borrowed("An argument for '[opt]' is required")),
 		}
 	}
 
@@ -78,6 +80,16 @@ impl Scurvy {
 			.flatten()
 	}
 
+	pub fn get_req(&self, key: &str) -> &str {
+		match self.get(key) {
+			None => {
+				let pair = self.get_pair(key).unwrap();
+				self.print_missing_and_die(pair.key.preferred_key());
+			}
+			Some(s) => s,
+		}
+	}
+
 	pub fn parse<T: FromStr, F: Into<Formula<T>>>(&self, key: &str, formula: F) -> Option<T> {
 		let formula = formula.into();
 		let Some(got) = self.get(key) else {
@@ -115,8 +127,7 @@ impl Scurvy {
 			None => match missing {
 				None => {
 					let pair = self.get_pair(key).unwrap();
-					eprintln!("An argument for '{}' is required", pair.key.preferred_key());
-					std::process::exit(-1);
+					self.print_missing_and_die(pair.key.preferred_key());
 				}
 				Some(misstr) => {
 					let pair = self.get_pair(key).unwrap();
@@ -129,6 +140,18 @@ impl Scurvy {
 			Some(o) => o,
 		}
 	}
+
+	pub fn set_missing_string<S: Into<Cow<'static, str>>>(&self, msg: S) {
+		self.missing_string.set(msg.into());
+	}
+
+	fn print_missing_and_die(&self, key: &str) -> ! {
+		// We can straight up take() here because of how we use this string,
+		// because we exit right after it
+		let str = self.missing_string.take().replace("[opt]", key);
+		eprintln!("{str}");
+		std::process::exit(-1);
+	}
 }
 
 pub struct Argument {