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        #[automatically_derived]
91        impl #impl_generics #strum_module_path::EnumProperty for #name #ty_generics #where_clause {
92            #[inline]
93            fn get_str(&self, prop: &str) -> ::core::option::Option<&'static str> {
94                match self {
95                    #(#built_string_arms),*
96                }
97            }
98
99            #[inline]
100            fn get_int(&self, prop: &str) -> ::core::option::Option<i64> {
101                match self {
102                    #(#built_int_arms),*
103                }
104            }
105
106            #[inline]
107            fn get_bool(&self, prop: &str) -> ::core::option::Option<bool> {
108                match self {
109                    #(#built_bool_arms),*
110                }
111            }
112
113        }
114    })
115}