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}