enum-status-code-macros 0.3.1

extern crate proc_macro;

use std::mem::take;

use itertools::Itertools;
use manyhow::manyhow;
use proc_macro2::TokenStream;
use quote::{ToTokens, quote};
use syn::{Attribute, Error, Expr, ItemEnum, TypeGenerics, Variant, parse2, spanned::Spanned};

#[manyhow(proc_macro_attribute)]
pub fn status_code(attr: TokenStream, item: TokenStream) -> syn::Result<TokenStream> {
    // Parse the enum and get some data about it
    let mut item: ItemEnum = parse2(item)?;
    let enum_name = &item.ident;
    let (impl_generics, ty_generics, where_clause) = item.generics.split_for_impl();
    let enum_span = item.span();

    // Get the default attribute, if any
    let default = (!attr.is_empty()).then(|| parse2(attr));

    item.variants
        // Iterate through all the variants
        .iter_mut()
        // Parse the variants
        .map(parse_variant(default, enum_name, &ty_generics))
        // Extract the error and reduce the variants into the body of a match
        .process_results(|iter| {
            iter.reduce(|a, b| {
                quote! {
                    #a
                    #b
                }
            })
        })?
        // Wrap it with the match syntax
        .map(|contents| {
            quote! {
                #item

                impl #impl_generics enum_status_code::AsStatusCode for #enum_name #ty_generics #where_clause {
                    fn as_status_code(&self) -> enum_status_code::http::StatusCode {
                        match self {
                            #contents
                        }
                    }
                }
            }
        })
        // Complain if there are no variants
        .ok_or(Error::new(enum_span, "expcted enum body"))
}

fn parse_variant(
    default: Option<Result<Expr, Error>>,
    enum_name: impl ToTokens,
    ty_generics: &TypeGenerics,
) -> impl Fn(&mut Variant) -> Result<TokenStream, Error> {
    move |variant| {
        // Get the pattern attribute and code attribute
        let (pattern_attribute, code_attribute) = get_attributes(variant)?;

        code_attribute
            .as_ref()
            // Parse the code attribute
            .map(Attribute::parse_args::<Expr>)
            // If the code attribute is none, use default
            .or(default.clone())
            // If both are none complain
            .ok_or(Error::new(
                variant.span(),
                "variant is missing status code and no default set",
            ))
            // Parse the attribute
            .map(parse_attribute(
                variant,
                ty_generics,
                &pattern_attribute,
                &enum_name,
            ))?
    }
}

fn get_attributes(variant: &mut Variant) -> Result<(Option<Attribute>, Option<Attribute>), Error> {
    // Both attributes are initially none
    let mut attributes = (None, None);
    // Iterate through all the attributes
    for attribute in take(&mut variant.attrs) {
        // Get the attribute "path" (#[this_bit(...)])
        let path = attribute.path();
        // Replace pattern with this if it's a pattern, and complain if anything is swapped out
        if path.is_ident("pattern") {
            if let Some(replaced) = attributes.0.replace(attribute) {
                return Err(Error::new(replaced.span(), "duplicate pattern attribute"));
            }
        } else if path.is_ident("code") {
            if let Some(replaced) = attributes.1.replace(attribute) {
                return Err(Error::new(replaced.span(), "duplicate code attribute"));
            }
        }
        // If both are filled, stop bothering
        if attributes.0.is_some() && attributes.1.is_some() {
            break;
        };
    }
    Ok(attributes)
}

fn parse_attribute(
    variant: &Variant,
    ty_generics: &TypeGenerics,
    pattern_attribute: &Option<Attribute>,
    enum_name: impl ToTokens,
) -> impl Fn(Result<Expr, Error>) -> Result<TokenStream, Error> {
    move |value| {
        // Some stuff we need
        let variant_name = &variant.ident;
        let expression = value?.into_token_stream();
        let turbofish = ty_generics.as_turbofish();

        // If the variant isn't a unit variant
        let pattern = if !matches!(variant.fields, syn::Fields::Unit)
            // And there is some pattern attribute
            && let Some(pattern) = pattern_attribute
        {
            // Return the args as a TokenStream
            Some(pattern.parse_args::<TokenStream>()?)
        } else {
            None
        };
        // Output the match arm
        Ok(quote! {
            #enum_name #turbofish :: #variant_name #pattern => #expression,
        })
    }
}