strum_macros/macros/
enum_properties.rs

1use std::collections::HashMap;
2
3use proc_macro2::TokenStream;
4use quote::quote;
5use syn::{Data, DeriveInput, Fields, Lit};
6
7use crate::helpers::{non_enum_error, HasStrumVariantProperties, HasTypeProperties};
8
9#[derive(Hash, PartialEq, Eq)]
10enum PropertyType {
11    String,
12    Integer,
13    Bool,
14}
15
16const PROPERTY_TYPES: [PropertyType; 3] = [
17    PropertyType::String,
18    PropertyType::Integer,
19    PropertyType::Bool,
20];
21
22pub fn enum_properties_inner(ast: &DeriveInput) -> syn::Result<TokenStream> {
23    let name = &ast.ident;
24    let (impl_generics, ty_generics, where_clause) = ast.generics.split_for_impl();
25    let variants = match &ast.data {
26        Data::Enum(v) => &v.variants,
27        _ => return Err(non_enum_error()),
28    };
29    let type_properties = ast.get_type_properties()?;
30    let strum_module_path = type_properties.crate_module_path();
31
32    let mut built_arms: HashMap<_, _> = PROPERTY_TYPES.iter().map(|p| (p, Vec::new())).collect();
33
34    for variant in variants {
35        let ident = &variant.ident;
36        let variant_properties = variant.get_variant_properties()?;
37        let mut arms: HashMap<_, _> = PROPERTY_TYPES.iter().map(|p| (p, Vec::new())).collect();
38        // But you can disable the messages.
39        if variant_properties.disabled.is_some() {
40            continue;
41        }
42
43        let params = match variant.fields {
44            Fields::Unit => quote! {},
45            Fields::Unnamed(..) => quote! { (..) },
46            Fields::Named(..) => quote! { {..} },
47        };
48
49        for (key, value) in variant_properties.props {
50            let property_type = match value {
51                Lit::Str(..) => PropertyType::String,
52                Lit::Bool(..) => PropertyType::Bool,
53                Lit::Int(..) => PropertyType::Integer,
54                _ => todo!("TODO"),
55            };
56
57            arms.get_mut(&property_type)
58                .unwrap()
59                .push(quote! { #key => ::core::option::Option::Some( #value )});
60        }
61
62        for property in &PROPERTY_TYPES {
63            arms.get_mut(&property)
64                .unwrap()
65                .push(quote! { _ => ::core::option::Option::None });
66            let arms_as_string = &arms[property];
67            built_arms.get_mut(&property).unwrap().push(quote! {
68                &#name::#ident #params => {
69                    match prop {
70                        #(#arms_as_string),*
71                    }
72                }
73            });
74        }
75    }
76
77    for (_, arms) in built_arms.iter_mut() {
78        if arms.len() < variants.len() {
79            arms.push(quote! { _ => ::core::option::Option::None });
80        }
81    }
82
83    let (built_string_arms, built_int_arms, built_bool_arms) = (
84        &built_arms[&PropertyType::String],
85        &built_arms[&PropertyType::Integer],
86        &built_arms[&PropertyType::Bool],
87    );
88
89    Ok(quote! {
90        impl #impl_generics #strum_module_path::EnumProperty for #name #ty_generics #where_clause {
91            #[inline]
92            fn get_str(&self, prop: &str) -> ::core::option::Option<&'static str> {
93                match self {
94                    #(#built_string_arms),*
95                }
96            }
97
98            #[inline]
99            fn get_int(&self, prop: &str) -> ::core::option::Option<i64> {
100                match self {
101                    #(#built_int_arms),*
102                }
103            }
104
105            #[inline]
106            fn get_bool(&self, prop: &str) -> ::core::option::Option<bool> {
107                match self {
108                    #(#built_bool_arms),*
109                }
110            }
111
112        }
113    })
114}