use proc_macro::TokenStream;
use quote::quote;
use syn::{Attribute, Data, DeriveInput, Error, Expr, Fields, Ident, Variant, parse_macro_input};
#[proc_macro_derive(AsStatusCode, attributes(as_status_code))]
pub fn as_status_code(input: TokenStream) -> TokenStream {
let input = parse_macro_input!(input as DeriveInput);
let name = &input.ident;
match &input.data {
Data::Struct(_) => struct_as_status_code(name, &input.attrs),
Data::Enum(data) => enum_as_status_code(name, data.variants.iter()),
Data::Union(_) => Error::new(name.span(), "union is not supported")
.to_compile_error()
.into(),
}
}
fn struct_as_status_code(name: &Ident, struct_attrs: &[Attribute]) -> TokenStream {
let Some(attr) = struct_attrs
.iter()
.find(|attr| attr.path().is_ident("as_status_code"))
else {
return Error::new(name.span(), "missing #[as_status_code(...)] attribute")
.to_compile_error()
.into();
};
let expr = match attr.parse_args::<Expr>() {
Ok(expr) => expr,
Err(err) => return err.to_compile_error().into(),
};
let expanded = quote! {
impl ::as_http_status_code::AsStatusCode for #name {
fn as_status_code(&self) -> ::as_http_status_code::StatusCode {
#expr
}
}
};
expanded.into()
}
fn enum_as_status_code<'a>(
name: &Ident,
variants: impl Iterator<Item = &'a Variant>,
) -> TokenStream {
let mut match_arms = Vec::new();
for variant in variants {
let ident = &variant.ident;
let Some(attr) = variant
.attrs
.iter()
.find(|attr| attr.path().is_ident("as_status_code"))
else {
return Error::new(ident.span(), "missing #[as_status_code(...)] attribute")
.to_compile_error()
.into();
};
let expr = match attr.parse_args::<Expr>() {
Ok(expr) => expr,
Err(err) => return err.to_compile_error().into(),
};
let quote = match &variant.fields {
Fields::Named(_) => quote! {
Self::#ident { .. } => #expr,
},
Fields::Unnamed(_) => quote! {
Self::#ident(..) => #expr,
},
Fields::Unit => quote! {
Self::#ident => #expr,
},
};
match_arms.push(quote);
}
let expanded = quote! {
impl ::as_http_status_code::AsStatusCode for #name {
fn as_status_code(&self) -> ::as_http_status_code::StatusCode {
match self {
#(#match_arms)*
}
}
}
};
expanded.into()
}