use darling::FromMeta;
use proc_macro_error2::abort;
use syn::spanned::Spanned;
use syn::{Attribute, ItemFn};
pub(crate) enum FuncAttribute {
Init,
Test(TestAttribute),
ShouldPanic,
Ignore,
Timeout(TimeoutAttribute),
}
impl FuncAttribute {
fn try_from_attr(attr: &Attribute) -> Option<Self> {
let ident = attr.path().get_ident()?.to_string();
Some(match ident.as_str() {
"init" => FuncAttribute::Init,
"test" => FuncAttribute::Test(TestAttribute::from_attr(attr)),
"should_panic" => FuncAttribute::ShouldPanic,
"ignore" => FuncAttribute::Ignore,
"timeout" => FuncAttribute::Timeout(TimeoutAttribute::from_attr(attr)),
_ => return None,
})
}
}
pub(crate) struct TimeoutAttribute {
pub value: u32,
}
impl syn::parse::Parse for TimeoutAttribute {
fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
let value_lit: syn::LitInt = input.parse()?;
let value = value_lit.base10_parse::<u32>()?;
Ok(TimeoutAttribute { value })
}
}
impl TimeoutAttribute {
fn from_attr(attr: &Attribute) -> Self {
match attr.parse_args::<TimeoutAttribute>() {
Ok(timeout_attr) => timeout_attr,
Err(e) => {
abort!(
attr,
"failed to parse `timeout` attribute. Must be of the form #[timeout(10)] where 10 is the timeout in seconds. Error: {}",
e
);
}
}
}
}
#[derive(Debug, FromMeta, Default)]
pub(crate) struct TestAttribute {
#[darling(default)]
pub init: Option<syn::Ident>,
}
impl TestAttribute {
pub fn from_attr(attr: &Attribute) -> Self {
match &attr.meta {
syn::Meta::Path(_) => TestAttribute::default(),
meta => match TestAttribute::from_meta(meta) {
Ok(test_attr) => test_attr,
Err(e) => abort!(attr, "failed to parse `test` attribute. Must be of the form #[test(init = init_function)]: {}", e),
},
}
}
}
pub(crate) struct FunctionWithAttributes {
pub func: ItemFn,
pub attributes: Vec<(FuncAttribute, proc_macro2::Span)>,
}
impl From<ItemFn> for FunctionWithAttributes {
fn from(mut func: ItemFn) -> Self {
let mut attributes = vec![];
func.attrs.retain(|attr| {
if let Some(func_attr) = FuncAttribute::try_from_attr(attr) {
attributes.push((func_attr, attr.path().span()));
false
} else {
true
}
});
Self { func, attributes }
}
}