as_http_status_code_derive 0.1.0

A derive macro for defining HTTP status code representation for struct and enum variants
Documentation
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;

        // Look for the `#[as_status_code(...)]` attribute
        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()
}