From cd38fbd58d047bb83cc8217c4ab9998a4e4baa59 Mon Sep 17 00:00:00 2001 From: Aiden McClelland Date: Mon, 12 Apr 2021 12:07:16 -0600 Subject: [PATCH] add `count` and `multiple` --- .../src/command/build.rs | 39 ++++++--- .../src/command/mod.rs | 4 +- .../src/command/parse.rs | 81 ++++++++++++++++--- rpc-toolkit-macro/src/lib.rs | 2 + 4 files changed, 102 insertions(+), 24 deletions(-) diff --git a/rpc-toolkit-macro-internals/src/command/build.rs b/rpc-toolkit-macro-internals/src/command/build.rs index 7f91329..c7a4933 100644 --- a/rpc-toolkit-macro-internals/src/command/build.rs +++ b/rpc-toolkit-macro-internals/src/command/build.rs @@ -31,7 +31,7 @@ fn build_app(name: LitStr, opt: &mut Options, params: &mut [ParamType]) -> Token .iter_mut() .filter_map(|param| { if let ParamType::Arg(arg) = param { - if arg.stdin { + if arg.stdin.is_some() { return None; } let name = arg.name.clone().unwrap(); @@ -50,19 +50,28 @@ fn build_app(name: LitStr, opt: &mut Options, params: &mut [ParamType]) -> Token modifications.extend(quote_spanned! { ty_span => arg = arg.takes_value(false); }); + } else if arg.count.is_some() { + modifications.extend(quote_spanned! { ty_span => + arg = arg.takes_value(false); + arg = arg.multiple(true); + }); } else { modifications.extend(quote_spanned! { ty_span => arg = arg.takes_value(true); }); - if p.path.segments.last().unwrap().ident != "Option" { - modifications.extend(quote_spanned! { ty_span => - arg = arg.required(true); - }); - } else { + if p.path.segments.last().unwrap().ident == "Option" { arg.optional = true; modifications.extend(quote_spanned! { ty_span => arg = arg.required(false); }); + } else if arg.multiple.is_some() { + modifications.extend(quote_spanned! { ty_span => + arg = arg.multiple(true); + }); + } else { + modifications.extend(quote_spanned! { ty_span => + arg = arg.required(true); + }); } } }; @@ -483,7 +492,7 @@ fn cli_handler( let name = arg.name.clone().unwrap(); let arg_name = LitStr::new(&name.to_string(), name.span()); let field_name = Ident::new(&format!("arg_{}", name), name.span()); - if arg.stdin { + if arg.stdin.is_some() { quote! { #field_name: rpc_toolkit_prelude::default_stdin_parser(&mut std::io::stdin(), matches)?, } @@ -491,29 +500,37 @@ fn cli_handler( quote! { #field_name: matches.is_present(#arg_name), } + } else if arg.count.is_some() { + quote! { + #field_name: matches.occurrences_of(#arg_name), + } } else { let parse_val = if let Some(parse) = &arg.parse { quote! { - #parse(arg_val, matches)? + #parse(arg_val, matches) } } else { quote! { - rpc_toolkit_prelude::default_arg_parser(arg_val, matches)? + rpc_toolkit_prelude::default_arg_parser(arg_val, matches) } }; if arg.optional { quote! { #field_name: if let Some(arg_val) = matches.value_of(#arg_name) { - Some(#parse_val) + Some(#parse_val?) } else { None }, } + } else if arg.multiple.is_some() { + quote! { + #field_name: matches.values_of(#arg_name).iter().flatten().map(|arg_val| #parse_val).collect::>()?, + } } else { quote! { #field_name: { let arg_val = matches.value_of(#arg_name).unwrap(); - #parse_val + #parse_val? }, } } diff --git a/rpc-toolkit-macro-internals/src/command/mod.rs b/rpc-toolkit-macro-internals/src/command/mod.rs index 73c7666..19348b8 100644 --- a/rpc-toolkit-macro-internals/src/command/mod.rs +++ b/rpc-toolkit-macro-internals/src/command/mod.rs @@ -76,7 +76,9 @@ pub struct ArgOptions { short: Option, long: Option, parse: Option, - stdin: bool, + count: Option, + multiple: Option, + stdin: Option, } pub enum ParamType { diff --git a/rpc-toolkit-macro-internals/src/command/parse.rs b/rpc-toolkit-macro-internals/src/command/parse.rs index 4da703c..17f2281 100644 --- a/rpc-toolkit-macro-internals/src/command/parse.rs +++ b/rpc-toolkit-macro-internals/src/command/parse.rs @@ -317,7 +317,7 @@ pub fn parse_command_attr(args: AttributeArgs) -> Result { Ok(opt) } -pub fn parse_arg_attr(attr: Attribute, arg: PatType, has_stdin: &mut bool) -> Result { +pub fn parse_arg_attr(attr: Attribute, arg: PatType) -> Result { let arg_span = arg.span(); let mut opt = ArgOptions { ty: *arg.ty, @@ -331,7 +331,9 @@ pub fn parse_arg_attr(attr: Attribute, arg: PatType, has_stdin: &mut bool) -> Re short: None, long: None, parse: None, - stdin: false, + count: None, + multiple: None, + stdin: None, }; match attr.parse_meta()? { Meta::List(list) => { @@ -370,9 +372,6 @@ pub fn parse_arg_attr(attr: Attribute, arg: PatType, has_stdin: &mut bool) -> Re return Err(Error::new(nv.path.span(), "`parse` cannot be assigned to")); } NestedMeta::Meta(Meta::Path(p)) if p.is_ident("stdin") => { - if *has_stdin { - return Err(Error::new(p.span(), "duplicate argument `stdin`")); - } if opt.short.is_some() { return Err(Error::new( p.span(), @@ -391,8 +390,19 @@ pub fn parse_arg_attr(attr: Attribute, arg: PatType, has_stdin: &mut bool) -> Re "`stdin` and `help` are mutually exclusive", )); } - opt.stdin = true; - *has_stdin = true; + if opt.count.is_some() { + return Err(Error::new( + p.span(), + "`stdin` and `count` are mutually exclusive", + )); + } + if opt.multiple.is_some() { + return Err(Error::new( + p.span(), + "`stdin` and `multiple` are mutually exclusive", + )); + } + opt.stdin = Some(p); } NestedMeta::Meta(Meta::List(list)) if list.path.is_ident("stdin") => { return Err(Error::new( @@ -403,12 +413,60 @@ pub fn parse_arg_attr(attr: Attribute, arg: PatType, has_stdin: &mut bool) -> Re NestedMeta::Meta(Meta::NameValue(nv)) if nv.path.is_ident("stdin") => { return Err(Error::new(nv.path.span(), "`stdin` cannot be assigned to")); } + NestedMeta::Meta(Meta::Path(p)) if p.is_ident("count") => { + if opt.stdin.is_some() { + return Err(Error::new( + p.span(), + "`stdin` and `count` are mutually exclusive", + )); + } + if opt.multiple.is_some() { + return Err(Error::new( + p.span(), + "`count` and `multiple` are mutually exclusive", + )); + } + opt.count = Some(p); + } + NestedMeta::Meta(Meta::List(list)) if list.path.is_ident("count") => { + return Err(Error::new( + list.path.span(), + "`count` does not take any arguments", + )); + } + NestedMeta::Meta(Meta::NameValue(nv)) if nv.path.is_ident("count") => { + return Err(Error::new(nv.path.span(), "`count` cannot be assigned to")); + } + NestedMeta::Meta(Meta::Path(p)) if p.is_ident("multiple") => { + if opt.stdin.is_some() { + return Err(Error::new( + p.span(), + "`stdin` and `multiple` are mutually exclusive", + )); + } + if opt.count.is_some() { + return Err(Error::new( + p.span(), + "`count` and `multiple` are mutually exclusive", + )); + } + opt.multiple = Some(p); + } + NestedMeta::Meta(Meta::List(list)) if list.path.is_ident("multiple") => { + return Err(Error::new( + list.path.span(), + "`multiple` does not take any arguments", + )); + } + NestedMeta::Meta(Meta::NameValue(nv)) if nv.path.is_ident("count") => { + return Err(Error::new(nv.path.span(), "`count` cannot be assigned to")); + } NestedMeta::Meta(Meta::NameValue(nv)) if nv.path.is_ident("help") => { if let Lit::Str(help) = nv.lit { if opt.help.is_some() { return Err(Error::new(help.span(), "duplicate argument `help`")); } - if opt.stdin { + if opt.stdin.is_some() { return Err(Error::new( help.span(), "`stdin` and `help` are mutually exclusive", @@ -464,7 +522,7 @@ pub fn parse_arg_attr(attr: Attribute, arg: PatType, has_stdin: &mut bool) -> Re if opt.short.is_some() { return Err(Error::new(short.span(), "duplicate argument `short`")); } - if opt.stdin { + if opt.stdin.is_some() { return Err(Error::new( short.span(), "`stdin` and `short` are mutually exclusive", @@ -489,7 +547,7 @@ pub fn parse_arg_attr(attr: Attribute, arg: PatType, has_stdin: &mut bool) -> Re if opt.long.is_some() { return Err(Error::new(long.span(), "duplicate argument `long`")); } - if opt.stdin { + if opt.stdin.is_some() { return Err(Error::new( long.span(), "`stdin` and `long` are mutually exclusive", @@ -529,7 +587,6 @@ pub fn parse_arg_attr(attr: Attribute, arg: PatType, has_stdin: &mut bool) -> Re pub fn parse_param_attrs(item: &mut ItemFn) -> Result> { let mut params = Vec::new(); - let mut has_stdin = false; for param in item.sig.inputs.iter_mut() { if let FnArg::Typed(param) = param { let mut ty = ParamType::None; @@ -538,7 +595,7 @@ pub fn parse_param_attrs(item: &mut ItemFn) -> Result> { if param.attrs[i].path.is_ident("arg") { let attr = param.attrs.remove(i); if matches!(ty, ParamType::None) { - ty = ParamType::Arg(parse_arg_attr(attr, param.clone(), &mut has_stdin)?); + ty = ParamType::Arg(parse_arg_attr(attr, param.clone())?); } else if matches!(ty, ParamType::Arg(_)) { return Err(Error::new( attr.span(), diff --git a/rpc-toolkit-macro/src/lib.rs b/rpc-toolkit-macro/src/lib.rs index 803260c..42b3b0c 100644 --- a/rpc-toolkit-macro/src/lib.rs +++ b/rpc-toolkit-macro/src/lib.rs @@ -18,6 +18,8 @@ pub fn command(args: TokenStream, item: TokenStream) -> TokenStream { /// - note: `arg` is the type of the argument /// - `#[arg(stdin)]` -> Parse the argument from stdin when using the CLI /// - `custom_parse_fn :: Into err => (&[u8], &ArgMatches<'_>) -> Result` +/// - `#[arg(count)]` -> Treat argument as flag, count occurrences +/// - `#[arg(multiple)]` -> Allow the arg to be specified multiple times. Collect the args after parsing. #[proc_macro_attribute] pub fn arg(_: TokenStream, _: TokenStream) -> TokenStream { syn::Error::new(