tonic_thiserror_impl/
lib.rs

1extern crate proc_macro;
2
3use proc_macro::TokenStream;
4use quote::quote;
5use syn::{parse_macro_input, DeriveInput, Attribute, Meta, NestedMeta, Data};
6
7#[proc_macro_derive(TonicThisError, attributes(code))]
8pub fn tonic_thiserror_derive(input: TokenStream) -> TokenStream {
9    let ast = parse_macro_input!(input as DeriveInput);
10    impl_tonic_thiserror(&ast).unwrap_or_else(|e| e.to_compile_error()).into()
11}
12
13fn impl_tonic_thiserror(ast: &DeriveInput) -> syn::Result<proc_macro2::TokenStream> {
14    let name = &ast.ident;
15    let variants = match &ast.data {
16        Data::Enum(data) => &data.variants,
17        _ => return Err(syn::Error::new_spanned(&ast, "TonicThisError can only be used on enums")),
18    };
19
20    let mut arms = proc_macro2::TokenStream::new();
21
22    for variant in variants {
23        let ident = &variant.ident;
24        let code_attr = variant
25            .attrs
26            .iter()
27            .find(|attr| attr.path.is_ident("code"))
28            .ok_or_else(|| syn::Error::new_spanned(variant, "Missing `code` attribute"))?;
29
30        let code = parse_code_attr(code_attr)?;
31
32        let arm = quote! {
33            #name::#ident { .. } => ::tonic::Status::new(#code, format!("{}", err)),
34        };
35
36        arms.extend(arm);
37    }
38
39    let gen = quote! {
40        impl From<#name> for ::tonic::Status {
41            fn from(err: #name) -> ::tonic::Status {
42                match err {
43                    #arms
44                }
45            }
46        }
47    };
48
49    Ok(gen)
50}
51
52fn parse_code_attr(attr: &Attribute) -> syn::Result<proc_macro2::TokenStream> {
53    let meta = attr.parse_meta()?;
54    let nested = match meta {
55        Meta::List(meta) => {
56            if meta.nested.len() != 1 {
57                return Err(syn::Error::new_spanned(attr, "Expected exactly one `code`"));
58            }
59            meta.nested.first().unwrap().clone()
60        }
61        _ => return Err(syn::Error::new_spanned(attr, "Expected `code` attribute to be a list")),
62    };
63
64    match nested {
65        NestedMeta::Meta(Meta::Path(path)) => {
66            // Assuming the code is one of tonic::Code variants
67            // No need to parse a string; it's directly the variant name
68            if path.segments.len() == 1 {
69                let segment = &path.segments[0];
70                let ident = &segment.ident;
71                Ok(quote! { ::tonic::Code::#ident })
72            } else {
73                Err(syn::Error::new_spanned(path, "Expected tonic::Code variant"))
74            }
75        }
76        _ => Err(syn::Error::new_spanned(attr, "Expected tonic::Code variant for `code` attribute")),
77    }
78}