enum_status_code_macros/
lib.rs

1extern crate proc_macro;
2
3use std::mem::take;
4
5use itertools::Itertools;
6use manyhow::manyhow;
7use proc_macro2::TokenStream;
8use quote::{ToTokens, quote};
9use syn::{Attribute, Error, Expr, ItemEnum, TypeGenerics, Variant, parse2, spanned::Spanned};
10
11#[manyhow(proc_macro_attribute)]
12pub fn status_code(attr: TokenStream, item: TokenStream) -> syn::Result<TokenStream> {
13    // Parse the enum and get some data about it
14    let mut item: ItemEnum = parse2(item)?;
15    let enum_name = &item.ident;
16    let (impl_generics, ty_generics, where_clause) = item.generics.split_for_impl();
17    let enum_span = item.span();
18
19    // Get the default attribute, if any
20    let default = (!attr.is_empty()).then(|| parse2(attr));
21
22    item.variants
23        // Iterate through all the variants
24        .iter_mut()
25        // Parse the variants
26        .map(parse_variant(default, enum_name, &ty_generics))
27        // Extract the error and reduce the variants into the body of a match
28        .process_results(|iter| {
29            iter.reduce(|a, b| {
30                quote! {
31                    #a
32                    #b
33                }
34            })
35        })?
36        // Wrap it with the match syntax
37        .map(|contents| {
38            quote! {
39                #item
40
41                impl #impl_generics enum_status_code::AsStatusCode for #enum_name #ty_generics #where_clause {
42                    fn as_status_code(&self) -> enum_status_code::http::StatusCode {
43                        match self {
44                            #contents
45                        }
46                    }
47                }
48            }
49        })
50        // Complain if there are no variants
51        .ok_or(Error::new(enum_span, "expcted enum body"))
52}
53
54fn parse_variant(
55    default: Option<Result<Expr, Error>>,
56    enum_name: impl ToTokens,
57    ty_generics: &TypeGenerics,
58) -> impl Fn(&mut Variant) -> Result<TokenStream, Error> {
59    move |variant| {
60        // Get the pattern attribute and code attribute
61        let (pattern_attribute, code_attribute) = get_attributes(variant)?;
62
63        code_attribute
64            .as_ref()
65            // Parse the code attribute
66            .map(Attribute::parse_args::<Expr>)
67            // If the code attribute is none, use default
68            .or(default.clone())
69            // If both are none complain
70            .ok_or(Error::new(
71                variant.span(),
72                "variant is missing status code and no default set",
73            ))
74            // Parse the attribute
75            .map(parse_attribute(
76                variant,
77                ty_generics,
78                &pattern_attribute,
79                &enum_name,
80            ))?
81    }
82}
83
84fn get_attributes(variant: &mut Variant) -> Result<(Option<Attribute>, Option<Attribute>), Error> {
85    // Both attributes are initially none
86    let mut attributes = (None, None);
87    // Iterate through all the attributes
88    for attribute in take(&mut variant.attrs) {
89        // Get the attribute "path" (#[this_bit(...)])
90        let path = attribute.path();
91        // Replace pattern with this if it's a pattern, and complain if anything is swapped out
92        if path.is_ident("pattern") {
93            if let Some(replaced) = attributes.0.replace(attribute) {
94                return Err(Error::new(replaced.span(), "duplicate pattern attribute"));
95            }
96        } else if path.is_ident("code") {
97            if let Some(replaced) = attributes.1.replace(attribute) {
98                return Err(Error::new(replaced.span(), "duplicate code attribute"));
99            }
100        }
101        // If both are filled, stop bothering
102        if attributes.0.is_some() && attributes.1.is_some() {
103            break;
104        };
105    }
106    Ok(attributes)
107}
108
109fn parse_attribute(
110    variant: &Variant,
111    ty_generics: &TypeGenerics,
112    pattern_attribute: &Option<Attribute>,
113    enum_name: impl ToTokens,
114) -> impl Fn(Result<Expr, Error>) -> Result<TokenStream, Error> {
115    move |value| {
116        // Some stuff we need
117        let variant_name = &variant.ident;
118        let expression = value?.into_token_stream();
119        let turbofish = ty_generics.as_turbofish();
120
121        // If the variant isn't a unit variant
122        let pattern = if !matches!(variant.fields, syn::Fields::Unit)
123            // And there is some pattern attribute
124            && let Some(pattern) = pattern_attribute
125        {
126            // Return the args as a TokenStream
127            Some(pattern.parse_args::<TokenStream>()?)
128        } else {
129            None
130        };
131        // Output the match arm
132        Ok(quote! {
133            #enum_name #turbofish :: #variant_name #pattern => #expression,
134        })
135    }
136}
enum_status_code_macros/lib.rs

enum_status_code_macros/
lib.rs
