use std::collections::BTreeMap;
#[derive(Debug)]
pub struct Error;
pub enum Value<'t> {
Bool(&'t mut bool),
Int8(&'t mut i8),
UInt8(&'t mut u8),
Int16(&'t mut i16),
UInt16(&'t mut u16),
Int32(&'t mut i32),
UInt32(&'t mut u32),
Int64(&'t mut i64),
UInt64(&'t mut u64),
Float32(&'t mut f32),
Float64(&'t mut f64),
String(&'t mut String),
}
pub trait ValueInner {
fn to_value<'t>(&'t mut self) -> Value<'t>;
}
impl ValueInner for bool {
fn to_value<'t>(&'t mut self) -> Value<'t> {
Value::Bool(self)
}
}
impl ValueInner for i64 {
fn to_value<'t>(&'t mut self) -> Value<'t> {
Value::Int64(self)
}
}
impl ValueInner for String {
fn to_value<'t>(&'t mut self) -> Value<'t> {
Value::String(self)
}
}
impl<'t> Value<'t> {
fn parse(&mut self, opt: &str) -> Result<(), Error> {
use Value as E;
match self {
E::Bool(x) => match opt {
"true" | "1" | "on" | "yes" => **x = true,
"false" | "0" | "off" | "no" => **x = false,
_ => return Err(Error),
},
E::Int8(x) => {
**x = opt.parse().map_err(|_| Error)?;
}
E::UInt8(x) => {
**x = opt.parse().map_err(|_| Error)?;
}
E::Int16(x) => {
**x = opt.parse().map_err(|_| Error)?;
}
E::UInt16(x) => {
**x = opt.parse().map_err(|_| Error)?;
}
E::Int32(x) => {
**x = opt.parse().map_err(|_| Error)?;
}
E::UInt32(x) => {
**x = opt.parse().map_err(|_| Error)?;
}
E::Int64(x) => {
**x = opt.parse().map_err(|_| Error)?;
}
E::UInt64(x) => {
**x = opt.parse().map_err(|_| Error)?;
}
E::Float32(x) => {
**x = opt.parse().map_err(|_| Error)?;
}
E::Float64(x) => {
**x = opt.parse().map_err(|_| Error)?;
}
E::String(x) => {
**x = opt.to_owned();
}
}
Ok(())
}
}
impl<'t> ToString for Value<'t> {
fn to_string(&self) -> String {
use Value as E;
match self {
E::Bool(x) => x.to_string(),
E::Int8(x) => x.to_string(),
E::UInt8(x) => x.to_string(),
E::Int16(x) => x.to_string(),
E::UInt16(x) => x.to_string(),
E::Int32(x) => x.to_string(),
E::UInt32(x) => x.to_string(),
E::Int64(x) => x.to_string(),
E::UInt64(x) => x.to_string(),
E::Float32(x) => x.to_string(),
E::Float64(x) => x.to_string(),
E::String(x) => format!("{:?}", x),
}
}
}
pub struct Flag<'t> {
name: String,
short: Option<char>,
value: Value<'t>,
usage: String,
defval: String,
}
impl<'t> Flag<'t> {
fn new(name: String, short: Option<char>, value: Value<'t>, usage: String) -> Flag<'t> {
let defval = value.to_string();
Flag {
name,
usage,
short,
value,
defval,
}
}
}
#[derive(Clone, Copy)]
pub enum ErrorHandler {
Exit,
Return,
}
pub struct FlagSet<'t> {
flags: Vec<Flag<'t>>,
name_map: BTreeMap<String, usize>,
short_name_map: BTreeMap<char, usize>,
prog: String,
error_handler: ErrorHandler,
}
impl<'t> FlagSet<'t> {
pub fn new_without_help() -> FlagSet<'t> {
FlagSet {
flags: Vec::new(),
name_map: BTreeMap::new(),
short_name_map: BTreeMap::new(),
prog: "<application name>".to_string(),
error_handler: ErrorHandler::Exit,
}
}
pub fn new(help: &'t mut bool) -> FlagSet<'t> {
let mut fs = Self::new_without_help();
fs.define("help", Some('h'), help, "Print this help and exit");
fs
}
pub fn define<Name: ToString, Usage: ToString, V: ValueInner>(
&mut self,
name: Name,
short: Option<char>,
b: &'t mut V,
usage: Usage,
) {
self.name_map.insert(name.to_string(), self.flags.len());
if let Some(short) = short {
self.short_name_map.insert(short, self.flags.len());
}
self.flags.push(Flag::new(
name.to_string(),
short,
b.to_value(),
usage.to_string(),
));
}
pub fn parse(&mut self, args: &[String]) -> Result<Vec<String>, Error> {
let eh = self.error_handler;
self.parse_inner(args).map_err(|err| {
if matches!(eh, ErrorHandler::Exit) {
let _ = self.usage(&mut std::io::stderr());
std::process::exit(1);
}
err
})
}
pub fn usage<W: std::io::Write>(&self, w: &mut W) -> std::io::Result<()> {
writeln!(w, "Usage: {} [flags]... [arguments]...", self.prog)?;
writeln!(w)?;
writeln!(w, "Flags:")?;
for flag in self.flags.iter() {
if let Some(short) = flag.short {
write!(w, "\t-{},--{}", short, flag.name)?;
} else {
write!(w, "\t--{}", flag.name)?;
}
writeln!(w, " (default: {})", flag.defval)?;
for line in flag.usage.split('\n') {
writeln!(w, "\t\t{}", line)?;
}
}
Ok(())
}
fn parse_inner(&mut self, args: &[String]) -> Result<Vec<String>, Error> {
self.prog = args[0].clone();
let mut args = &args[1..];
let mut pos = vec![];
while let Some((arg, rest)) = args.split_first() {
let Some(arg) = arg.strip_prefix('-') else {
pos.push(arg.to_owned());
args = rest;
continue;
};
if arg == "-" {
pos.extend(rest.to_vec());
break;
}
match arg.split_once('=') {
Some((name, opt)) => {
self.parse_with_eq(name, opt)?;
args = rest;
}
None => {
args = self.parse_without_eq(arg, rest)?;
}
}
}
Ok(pos)
}
fn parse_without_eq<'s>(
&mut self,
name: &str,
mut args: &'s [String],
) -> Result<&'s [String], Error> {
if let Some(long) = name.strip_prefix('-') {
let Some(flag) = self.name_map.get_mut(long) else {
return Err(Error);
};
let flag = &mut self.flags[*flag];
if matches!(flag.value, Value::Bool(_)) {
flag.value.parse("true")?;
} else {
let Some((next, rest)) = args.split_first() else {
return Err(Error);
};
flag.value.parse(next)?;
args = rest;
}
} else {
for short in name.chars() {
let Some(flag) = self.short_name_map.get_mut(&short) else {
return Err(Error);
};
self.flags[*flag].value.parse("true")?;
}
}
Ok(args)
}
fn parse_with_eq(&mut self, name: &str, opt: &str) -> Result<(), Error> {
let flag = if let Some(long) = name.strip_prefix('-') {
let Some(flag) = self.name_map.get_mut(long) else {
return Err(Error);
};
*flag
} else {
if name.len() != 1 {
return Err(Error);
}
let short = name.chars().next().unwrap();
let Some(flag) = self.short_name_map.get_mut(&short) else {
return Err(Error);
};
*flag
};
self.flags[flag].value.parse(opt)?;
Ok(())
}
}
#[test]
fn test_flag() {
#[derive(Debug, Default)]
struct TestCase {
args: Vec<String>,
create: bool,
gzip: bool,
strip_components: i64,
pos: Vec<String>,
}
let test_cases = [
TestCase {
args: vec!["tar".to_string(), "--create=true".to_string()],
create: true,
..Default::default()
},
TestCase {
args: vec!["tar".to_string(), "--create=false".to_string()],
..Default::default()
},
TestCase {
args: vec!["tar".to_string(), "-cz".to_string()],
create: true,
gzip: true,
..Default::default()
},
TestCase {
args: vec!["tar".to_string(), "--strip-components=1".to_string()],
strip_components: 1,
..Default::default()
},
TestCase {
args: vec!["tar".to_string(), "--strip-components=-10".to_string()],
strip_components: -10,
..Default::default()
},
TestCase {
args: vec![
"tar".to_string(),
"--strip-components=-10".to_string(),
"--strip-components=10".to_string(),
],
strip_components: 10,
..Default::default()
},
TestCase {
args: vec![
"tar".to_string(),
"--strip-components".to_string(),
"10".to_string(),
],
strip_components: 10,
..Default::default()
},
TestCase {
args: vec![
"tar".to_string(),
"input1.txt".to_string(),
"input2.txt".to_string(),
"-c".to_string(),
"input3.txt".to_string(),
"--".to_string(),
"-z".to_string(),
],
create: true,
pos: vec![
"input1.txt".to_string(),
"input2.txt".to_string(),
"input3.txt".to_string(),
"-z".to_string(),
],
..Default::default()
},
];
for test_case in test_cases {
let mut create = false;
let mut gzip = false;
let mut strip_comp = 0;
let mut fs = FlagSet::new_without_help();
fs.define(
"create",
Some('c'),
&mut create,
concat!(
"Create a new archive. Arguments supply the names of the files\n",
"to be archived. Directories are archived recursively, unless\n",
"the `--no-recursion' option is given."
),
);
fs.define(
"gzip",
Some('z'),
&mut gzip,
"Filter the archive through gzip(1).",
);
fs.define(
"strip-components",
None,
&mut strip_comp,
"Strip some leading components from file names on extraction.",
);
let pos = fs.parse(&test_case.args).unwrap();
assert_eq!(create, test_case.create);
assert_eq!(gzip, test_case.gzip);
assert_eq!(strip_comp, test_case.strip_components);
assert_eq!(pos, test_case.pos);
}
}