diff --git a/derive/src/lib.rs b/derive/src/lib.rs index c3ac1b6..56e0042 100644 --- a/derive/src/lib.rs +++ b/derive/src/lib.rs @@ -9,7 +9,7 @@ use syn::{ punctuated::Punctuated, Attribute, Data::Struct, - DeriveInput, Fields, LitStr, Token, + DeriveInput, Expr, Fields, Ident, LitStr, Token, }; #[derive(Eq, Hash, PartialEq, Debug)] @@ -24,11 +24,13 @@ enum OptionsAttribute { struct FlagAttribute { flags: Vec, + value: Option, } enum FlagArg { Short(char), Long(String), + Value(Expr), } // FIXME: Think of a better name @@ -70,10 +72,14 @@ pub fn options(input: TokenStream) -> TokenStream { } else { f.flags }; + + let stmt = match f.value { + Some(e) => quote!(self.#field_ident = #e;), + None => quote!(self.#field_ident = true;), + }; + for flag in flags { - map.entry(flag) - .or_default() - .push(quote!(self.#field_ident = true;)); + map.entry(flag).or_default().push(stmt.clone()); } } } @@ -119,7 +125,10 @@ fn parse_attr(attr: Attribute) -> Option { } fn parse_flag_attr(attr: Attribute) -> FlagAttribute { - let mut flag_attr = FlagAttribute { flags: vec![] }; + let mut flag_attr = FlagAttribute { + flags: vec![], + value: None, + }; let Ok(parsed_args) = attr .parse_args_with(Punctuated::::parse_terminated) else { @@ -129,6 +138,7 @@ fn parse_flag_attr(attr: Attribute) -> FlagAttribute { match arg { FlagArg::Long(s) => flag_attr.flags.push(Arg::Long(s)), FlagArg::Short(c) => flag_attr.flags.push(Arg::Short(c)), + FlagArg::Value(e) => flag_attr.value = Some(e), }; } flag_attr @@ -150,6 +160,15 @@ impl Parse for FlagArg { } panic!("Arguments to flag must start with \"-\" or \"--\""); } + + if input.peek(Ident) { + let name = input.parse::()?.to_string(); + input.parse::()?; + match name.as_str() { + "value" => return Ok(FlagArg::Value(input.parse::()?)), + _ => panic!("Unrecognized argument {} for flag attribute", name), + }; + } panic!("Arguments to flag attribute must be string literals"); } } diff --git a/tests/flags.rs b/tests/flags.rs index b0e4fc9..1cadf39 100644 --- a/tests/flags.rs +++ b/tests/flags.rs @@ -179,3 +179,79 @@ fn number_flag() { assert_eq!(Settings::parse(["-1"]).unwrap(), Settings { one: true }) } + +#[test] +fn false_bool() { + #[derive(Default, Options, PartialEq, Eq, Debug)] + struct Settings { + #[flag("-a")] + #[flag("-b", value = false)] + foo: bool, + } + + assert_eq!(Settings::parse(["-a"]).unwrap(), Settings { foo: true }); + assert_eq!(Settings::parse(["-b"]).unwrap(), Settings { foo: false }); + assert_eq!(Settings::parse(["-ab"]).unwrap(), Settings { foo: false }); + assert_eq!(Settings::parse(["-ba"]).unwrap(), Settings { foo: true }); + assert_eq!( + Settings::parse(["-a", "-b"]).unwrap(), + Settings { foo: false } + ); + assert_eq!( + Settings::parse(["-b", "-a"]).unwrap(), + Settings { foo: true } + ); +} + +#[test] +fn enum_flag() { + #[derive(Default, PartialEq, Eq, Debug)] + enum SomeEnum { + #[default] + VariantFoo, + VariantBar, + VariantBaz, + } + + #[derive(Default, Options, PartialEq, Eq, Debug)] + struct Settings { + #[flag(value = SomeEnum::VariantFoo)] + #[flag("--bar", value = SomeEnum::VariantBar)] + #[flag("--baz", value = SomeEnum::VariantBaz)] + foo: SomeEnum, + } + + assert_eq!( + Settings::parse(&[] as &[&str]).unwrap(), + Settings { + foo: SomeEnum::VariantFoo + } + ); + + assert_eq!( + Settings::parse(["--bar"]).unwrap(), + Settings { + foo: SomeEnum::VariantBar + } + ); + + assert_eq!( + Settings::parse(["--baz"]).unwrap(), + Settings { + foo: SomeEnum::VariantBaz + } + ); +} + +#[test] +fn count() { + #[derive(Default, Options)] + struct Settings { + #[flag(value = self.verbosity + 1)] + verbosity: u8, + } + + assert_eq!(Settings::parse(["-v"]).unwrap().verbosity, 1); + assert_eq!(Settings::parse(["-vv"]).unwrap().verbosity, 2); + assert_eq!(Settings::parse(["-vvv"]).unwrap().verbosity, 3); +}