add count and multiple

This commit is contained in:
Aiden McClelland
2021-04-12 12:07:16 -06:00
parent ceead74072
commit cd38fbd58d
4 changed files with 102 additions and 24 deletions

View File

@@ -31,7 +31,7 @@ fn build_app(name: LitStr, opt: &mut Options, params: &mut [ParamType]) -> Token
.iter_mut() .iter_mut()
.filter_map(|param| { .filter_map(|param| {
if let ParamType::Arg(arg) = param { if let ParamType::Arg(arg) = param {
if arg.stdin { if arg.stdin.is_some() {
return None; return None;
} }
let name = arg.name.clone().unwrap(); 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 => modifications.extend(quote_spanned! { ty_span =>
arg = arg.takes_value(false); 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 { } else {
modifications.extend(quote_spanned! { ty_span => modifications.extend(quote_spanned! { ty_span =>
arg = arg.takes_value(true); arg = arg.takes_value(true);
}); });
if p.path.segments.last().unwrap().ident != "Option" { if p.path.segments.last().unwrap().ident == "Option" {
modifications.extend(quote_spanned! { ty_span =>
arg = arg.required(true);
});
} else {
arg.optional = true; arg.optional = true;
modifications.extend(quote_spanned! { ty_span => modifications.extend(quote_spanned! { ty_span =>
arg = arg.required(false); 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 name = arg.name.clone().unwrap();
let arg_name = LitStr::new(&name.to_string(), name.span()); let arg_name = LitStr::new(&name.to_string(), name.span());
let field_name = Ident::new(&format!("arg_{}", name), name.span()); let field_name = Ident::new(&format!("arg_{}", name), name.span());
if arg.stdin { if arg.stdin.is_some() {
quote! { quote! {
#field_name: rpc_toolkit_prelude::default_stdin_parser(&mut std::io::stdin(), matches)?, #field_name: rpc_toolkit_prelude::default_stdin_parser(&mut std::io::stdin(), matches)?,
} }
@@ -491,29 +500,37 @@ fn cli_handler(
quote! { quote! {
#field_name: matches.is_present(#arg_name), #field_name: matches.is_present(#arg_name),
} }
} else if arg.count.is_some() {
quote! {
#field_name: matches.occurrences_of(#arg_name),
}
} else { } else {
let parse_val = if let Some(parse) = &arg.parse { let parse_val = if let Some(parse) = &arg.parse {
quote! { quote! {
#parse(arg_val, matches)? #parse(arg_val, matches)
} }
} else { } else {
quote! { quote! {
rpc_toolkit_prelude::default_arg_parser(arg_val, matches)? rpc_toolkit_prelude::default_arg_parser(arg_val, matches)
} }
}; };
if arg.optional { if arg.optional {
quote! { quote! {
#field_name: if let Some(arg_val) = matches.value_of(#arg_name) { #field_name: if let Some(arg_val) = matches.value_of(#arg_name) {
Some(#parse_val) Some(#parse_val?)
} else { } else {
None None
}, },
} }
} else if arg.multiple.is_some() {
quote! {
#field_name: matches.values_of(#arg_name).iter().flatten().map(|arg_val| #parse_val).collect::<Result<_, _>>()?,
}
} else { } else {
quote! { quote! {
#field_name: { #field_name: {
let arg_val = matches.value_of(#arg_name).unwrap(); let arg_val = matches.value_of(#arg_name).unwrap();
#parse_val #parse_val?
}, },
} }
} }

View File

@@ -76,7 +76,9 @@ pub struct ArgOptions {
short: Option<LitStr>, short: Option<LitStr>,
long: Option<LitStr>, long: Option<LitStr>,
parse: Option<Path>, parse: Option<Path>,
stdin: bool, count: Option<Path>,
multiple: Option<Path>,
stdin: Option<Path>,
} }
pub enum ParamType { pub enum ParamType {

View File

@@ -317,7 +317,7 @@ pub fn parse_command_attr(args: AttributeArgs) -> Result<Options> {
Ok(opt) Ok(opt)
} }
pub fn parse_arg_attr(attr: Attribute, arg: PatType, has_stdin: &mut bool) -> Result<ArgOptions> { pub fn parse_arg_attr(attr: Attribute, arg: PatType) -> Result<ArgOptions> {
let arg_span = arg.span(); let arg_span = arg.span();
let mut opt = ArgOptions { let mut opt = ArgOptions {
ty: *arg.ty, ty: *arg.ty,
@@ -331,7 +331,9 @@ pub fn parse_arg_attr(attr: Attribute, arg: PatType, has_stdin: &mut bool) -> Re
short: None, short: None,
long: None, long: None,
parse: None, parse: None,
stdin: false, count: None,
multiple: None,
stdin: None,
}; };
match attr.parse_meta()? { match attr.parse_meta()? {
Meta::List(list) => { 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")); return Err(Error::new(nv.path.span(), "`parse` cannot be assigned to"));
} }
NestedMeta::Meta(Meta::Path(p)) if p.is_ident("stdin") => { 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() { if opt.short.is_some() {
return Err(Error::new( return Err(Error::new(
p.span(), 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", "`stdin` and `help` are mutually exclusive",
)); ));
} }
opt.stdin = true; if opt.count.is_some() {
*has_stdin = true; 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") => { NestedMeta::Meta(Meta::List(list)) if list.path.is_ident("stdin") => {
return Err(Error::new( 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") => { NestedMeta::Meta(Meta::NameValue(nv)) if nv.path.is_ident("stdin") => {
return Err(Error::new(nv.path.span(), "`stdin` cannot be assigned to")); 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") => { NestedMeta::Meta(Meta::NameValue(nv)) if nv.path.is_ident("help") => {
if let Lit::Str(help) = nv.lit { if let Lit::Str(help) = nv.lit {
if opt.help.is_some() { if opt.help.is_some() {
return Err(Error::new(help.span(), "duplicate argument `help`")); return Err(Error::new(help.span(), "duplicate argument `help`"));
} }
if opt.stdin { if opt.stdin.is_some() {
return Err(Error::new( return Err(Error::new(
help.span(), help.span(),
"`stdin` and `help` are mutually exclusive", "`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() { if opt.short.is_some() {
return Err(Error::new(short.span(), "duplicate argument `short`")); return Err(Error::new(short.span(), "duplicate argument `short`"));
} }
if opt.stdin { if opt.stdin.is_some() {
return Err(Error::new( return Err(Error::new(
short.span(), short.span(),
"`stdin` and `short` are mutually exclusive", "`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() { if opt.long.is_some() {
return Err(Error::new(long.span(), "duplicate argument `long`")); return Err(Error::new(long.span(), "duplicate argument `long`"));
} }
if opt.stdin { if opt.stdin.is_some() {
return Err(Error::new( return Err(Error::new(
long.span(), long.span(),
"`stdin` and `long` are mutually exclusive", "`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<Vec<ParamType>> { pub fn parse_param_attrs(item: &mut ItemFn) -> Result<Vec<ParamType>> {
let mut params = Vec::new(); let mut params = Vec::new();
let mut has_stdin = false;
for param in item.sig.inputs.iter_mut() { for param in item.sig.inputs.iter_mut() {
if let FnArg::Typed(param) = param { if let FnArg::Typed(param) = param {
let mut ty = ParamType::None; let mut ty = ParamType::None;
@@ -538,7 +595,7 @@ pub fn parse_param_attrs(item: &mut ItemFn) -> Result<Vec<ParamType>> {
if param.attrs[i].path.is_ident("arg") { if param.attrs[i].path.is_ident("arg") {
let attr = param.attrs.remove(i); let attr = param.attrs.remove(i);
if matches!(ty, ParamType::None) { 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(_)) { } else if matches!(ty, ParamType::Arg(_)) {
return Err(Error::new( return Err(Error::new(
attr.span(), attr.span(),

View File

@@ -18,6 +18,8 @@ pub fn command(args: TokenStream, item: TokenStream) -> TokenStream {
/// - note: `arg` is the type of the argument /// - note: `arg` is the type of the argument
/// - `#[arg(stdin)]` -> Parse the argument from stdin when using the CLI /// - `#[arg(stdin)]` -> Parse the argument from stdin when using the CLI
/// - `custom_parse_fn :: Into<RpcError> err => (&[u8], &ArgMatches<'_>) -> Result<arg, err>` /// - `custom_parse_fn :: Into<RpcError> err => (&[u8], &ArgMatches<'_>) -> Result<arg, err>`
/// - `#[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] #[proc_macro_attribute]
pub fn arg(_: TokenStream, _: TokenStream) -> TokenStream { pub fn arg(_: TokenStream, _: TokenStream) -> TokenStream {
syn::Error::new( syn::Error::new(