Skip to main content

inquiry/
lib.rs

1use proc_macro::TokenStream;
2use quote::{format_ident, quote};
3use syn::{Data, DataEnum, DeriveInput, Variant, parse_macro_input};
4
5#[proc_macro_derive(Choice, attributes(desc))]
6pub fn derive_choice(input: TokenStream) -> TokenStream {
7    let input = parse_macro_input!(input as DeriveInput);
8    let enum_name = &input.ident;
9    let variants = match &input.data {
10        syn::Data::Enum(DataEnum { variants, .. }) => variants,
11        _ => panic!("Choice can only be derived for enums"),
12    };
13
14    let variant_matches = variants.iter().map(|variant| {
15        let variant_name = &variant.ident;
16        let desc = get_doc_comment(variant)
17            .unwrap_or_else(|| panic!("Missing description for variant {}", variant_name));
18
19        quote! {
20            #enum_name::#variant_name => #desc,
21        }
22    });
23
24    let variants = if let Data::Enum(data_enum) = &input.data {
25        data_enum
26            .variants
27            .iter()
28            .map(|v| &v.ident)
29            .collect::<Vec<_>>()
30    } else {
31        panic!("Choice only supports enums")
32    };
33
34    let module_name = format_ident!("__enum_choice_for_{}", enum_name.to_string().to_lowercase());
35
36    let expanded = quote! {
37        mod #module_name {
38            use super::*;
39
40            #[doc(hidden)]
41            pub trait Variants<T: 'static> {
42                const VARIANTS: &'static [T];
43            }
44
45            impl Variants<#enum_name> for #enum_name {
46                const VARIANTS: &'static [#enum_name] = &[#(#enum_name::#variants),*];
47            }
48
49            pub use Variants as VariantsTrait;
50        }
51
52        impl #enum_name {
53            fn description(&self) -> &'static str {
54                match self {
55                    #(#variant_matches)*
56                }
57            }
58
59            pub fn choice(msg: &str) -> ::inquire::error::InquireResult<Self>
60            where
61                Self: ::std::fmt::Display + ::std::fmt::Debug + Copy + Clone + 'static
62            {
63                let answer: Self = ::inquire::Select::new(msg, <Self as #module_name::VariantsTrait<Self>>::VARIANTS.to_vec())
64                    .prompt()?;
65
66                Ok(answer)
67            }
68        }
69
70        impl std::fmt::Display for #enum_name {
71            fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
72                write!(f, "{}", self.description())
73            }
74        }
75    };
76
77    TokenStream::from(expanded)
78}
79
80fn get_doc_comment(variant: &Variant) -> Option<String> {
81    let mut lines: Vec<_> = variant
82        .attrs
83        .iter()
84        .filter(|attr| attr.path().is_ident("doc"))
85        .filter_map(|attr| match &attr.meta {
86            syn::Meta::NameValue(syn::MetaNameValue {
87                value:
88                    syn::Expr::Lit(syn::ExprLit {
89                        lit: syn::Lit::Str(s),
90                        ..
91                    }),
92                ..
93            }) => Some(s.value()),
94            _ => None,
95        })
96        .skip_while(|s| s.trim().is_empty())
97        .flat_map(|s| {
98            s.split('\n')
99                .map(|s| {
100                    let s = s.strip_prefix(' ').unwrap_or(s);
101                    s.to_owned()
102                })
103                .collect::<Vec<_>>()
104        })
105        .collect();
106
107    while let Some(true) = lines.last().map(|s| s.trim().is_empty()) {
108        lines.pop();
109    }
110
111    Some(lines.concat())
112}