Allow passing iterations and seed as env variables

Co-Authored-By: Nathan Sobo <nathan@zed.dev>
This commit is contained in:
Antonio Scandurra 2021-07-28 17:48:51 +02:00 committed by Max Brunsfeld
parent 42e2b9ff4d
commit b0f3778381

View file

@ -13,7 +13,10 @@ pub fn test(args: TokenStream, function: TokenStream) -> TokenStream {
let args = syn::parse_macro_input!(args as AttributeArgs); let args = syn::parse_macro_input!(args as AttributeArgs);
let mut max_retries = 0; let mut max_retries = 0;
let mut num_iterations = 1; let mut num_iterations = 1;
let mut starting_seed = 0; let mut starting_seed = std::env::var("SEED")
.map(|i| i.parse().expect("invalid `SEED`"))
.ok();
for arg in args { for arg in args {
match arg { match arg {
NestedMeta::Meta(Meta::Path(name)) NestedMeta::Meta(Meta::Path(name))
@ -23,24 +26,34 @@ pub fn test(args: TokenStream, function: TokenStream) -> TokenStream {
} }
NestedMeta::Meta(Meta::NameValue(meta)) => { NestedMeta::Meta(Meta::NameValue(meta)) => {
let key_name = meta.path.get_ident().map(|i| i.to_string()); let key_name = meta.path.get_ident().map(|i| i.to_string());
let variable = match key_name.as_ref().map(String::as_str) { let result = (|| {
Some("retries") => &mut max_retries, match key_name.as_ref().map(String::as_str) {
Some("iterations") => &mut num_iterations, Some("retries") => max_retries = parse_int(&meta.lit)?,
Some("seed") => &mut starting_seed, Some("iterations") => {
_ => { if let Ok(iters) = std::env::var("ITERATIONS") {
return TokenStream::from( num_iterations = iters.parse().expect("invalid `ITERATIONS`");
syn::Error::new(meta.path.span(), "invalid argument") } else {
.into_compile_error(), num_iterations = parse_int(&meta.lit)?;
) }
}
Some("seed") => {
if starting_seed.is_none() {
starting_seed = Some(parse_int(&meta.lit)?);
}
}
_ => {
return Err(TokenStream::from(
syn::Error::new(meta.path.span(), "invalid argument")
.into_compile_error(),
))
}
} }
}; Ok(())
})();
*variable = match parse_int(&meta.lit) { if let Err(tokens) = result {
Ok(value) => value, return tokens;
Err(error) => { }
return TokenStream::from(error.into_compile_error());
}
};
} }
other => { other => {
return TokenStream::from( return TokenStream::from(
@ -49,6 +62,7 @@ pub fn test(args: TokenStream, function: TokenStream) -> TokenStream {
} }
} }
} }
let starting_seed = starting_seed.unwrap_or(0);
let mut inner_fn = parse_macro_input!(function as ItemFn); let mut inner_fn = parse_macro_input!(function as ItemFn);
let inner_fn_attributes = mem::take(&mut inner_fn.attrs); let inner_fn_attributes = mem::take(&mut inner_fn.attrs);
@ -142,10 +156,12 @@ pub fn test(args: TokenStream, function: TokenStream) -> TokenStream {
TokenStream::from(quote!(#outer_fn)) TokenStream::from(quote!(#outer_fn))
} }
fn parse_int(literal: &Lit) -> syn::Result<usize> { fn parse_int(literal: &Lit) -> Result<usize, TokenStream> {
if let Lit::Int(int) = &literal { let result = if let Lit::Int(int) = &literal {
int.base10_parse() int.base10_parse()
} else { } else {
Err(syn::Error::new(literal.span(), "must be an integer")) Err(syn::Error::new(literal.span(), "must be an integer"))
} };
result.map_err(|err| TokenStream::from(err.into_compile_error()))
} }