Skip to main content

enum_status_code_macros/
lib.rs

1extern crate proc_macro;
2
3use heck::AsSnekCase;
4use itertools::Itertools;
5use manyhow::manyhow;
6use proc_macro2::TokenStream;
7use quote::{ToTokens, quote};
8use syn::{Attribute, Error, Expr, Field, Ident, ItemEnum, parse2, spanned::Spanned};
9
10#[manyhow(proc_macro_attribute)]
11pub fn status_code(attr: TokenStream, item: TokenStream) -> syn::Result<TokenStream> {
12    let mut item: ItemEnum = parse2(item)?;
13    let enum_name = &item.ident;
14    let (impl_generics, ty_generics, where_clause) = item.generics.split_for_impl();
15    let enum_span = item.span();
16
17    let default = (!attr.is_empty()).then(|| parse2(attr));
18    item.variants
19        .iter_mut()
20        .map(|variant| {
21            let attribute = &variant
22                .attrs
23                .pop_if(|attr| attr.path().is_ident("code"));
24            let attribute = attribute.as_ref();
25            if let Some(value) = attribute
26                .map(Attribute::parse_args::<Expr>)
27                .or(default.clone())
28            {
29                let variant_name = &variant.ident;
30                let mut value = value?.into_token_stream();
31                let turbofish = ty_generics.as_turbofish();
32                let pattern: Option<TokenStream> = match &variant.fields {
33                    syn::Fields::Named(fields) => {
34                        fields.named.iter()
35                            .map(|Field { ident, .. }| ident.into_token_stream())
36                            .reduce(|a, b| quote! {
37                                #a, #b
38                            })
39                            .map(|p| quote! {
40                                { #p }
41                            })
42                    },
43                    syn::Fields::Unnamed(fields) => {
44                        let variant_name = Ident::new(&AsSnekCase(variant_name.to_string()).to_string(), enum_span);
45                        let tuple_members = (0..fields.unnamed.len()).map(|id| {
46                            let id_ident = Ident::new(&format!("__tuple_n{}", id), enum_span);
47                            id_ident.into_token_stream()
48                        }).reduce(|a, b| {
49                            quote! {
50                                #a, #b
51                            }
52                        });
53                        value = quote! {
54                            {
55                                let #variant_name = (#tuple_members);
56                                #value   
57                            }
58                        };
59                        Some(quote! {
60                            (#tuple_members)
61                        })
62                    },
63                    syn::Fields::Unit => None,
64                };
65                Ok(quote! {
66                    #enum_name #turbofish :: #variant_name #pattern => #value,
67                })
68            } else {
69                Err(Error::new(
70                    variant.span(),
71                    "variant is missing status code and no default set",
72                ))
73            }
74        })
75        .process_results(|iter| {
76            iter.tree_reduce(|a, b| {
77                quote! {
78                    #a
79                    #b
80                }
81            })
82        })?
83        .map(|contents| {
84            quote! {
85                #item
86
87                impl #impl_generics enum_status_code::AsStatusCode for #enum_name #ty_generics #where_clause {
88                    fn as_status_code(&self) -> enum_status_code::http::StatusCode {
89                        match self {
90                            #contents
91                        }
92                    }
93                }
94            }
95        })
96        .ok_or(Error::new(enum_span, "expcted enum body"))
97}