Repository: Peternator7/strum
Pull request #391 (2024-11-20 11:54:31)

Implement `get_int`/`get_bool` for properties
Fixes: #313

I am also suggesting to use `i64` as return type for `get_int` in order to cover also negative numbers.

Peternator7: Hey @marxin, thanks for the PR! This has been requested for quite a while. I tend to agree with @pali6 though that this would probably be better if accepted literals of the expected type rather than parsing the string.

In some ideal world, we could accept an expression `#[strum(props(my_prop = true || false))]` but that's difficult because we'd need to figure out the type. I think a better solution in the first version is to change the grammar so that `Prop` parses a [`Lit`](https://docs.rs/syn/latest/syn/enum.Lit.html) rather than a `LitStr` then we can do a pattern match in the macro to figure out what type of property we're looking at.
marxin: Here we go. I've implemented the suggested approach for the two functions. Eventually, I decided to track the 3 categories using HashMap in `enum_properties_inner`, but using arrays (the original commit in this PR) might be possible. Which approach do you prefer?

Patch:
diff --git a/strum/src/lib.rs b/strum/src/lib.rs
index 90f2d75..2a6d44c 100644
--- a/strum/src/lib.rs
+++ b/strum/src/lib.rs
@@ -155,9 +155,8 @@ pub trait EnumMessage {
 
 /// `EnumProperty` is a trait that makes it possible to store additional information
 /// with enum variants. This trait is designed to be used with the macro of the same
-/// name in the `strum_macros` crate. Currently, the only string literals are supported
-/// in attributes, the other methods will be implemented as additional attribute types
-/// become stabilized.
+/// name in the `strum_macros` crate. Currently, the string, integer and bool literals
+/// are supported in attributes.
 ///
 /// # Example
 ///
@@ -168,27 +167,24 @@ pub trait EnumMessage {
 ///
 /// #[derive(PartialEq, Eq, Debug, EnumProperty)]
 /// enum Class {
-///     #[strum(props(Teacher="Ms.Frizzle", Room="201"))]
+///     #[strum(props(Teacher="Ms.Frizzle", Room="201", students=16, mandatory=true))]
 ///     History,
 ///     #[strum(props(Teacher="Mr.Smith"))]
-///     #[strum(props(Room="103"))]
+///     #[strum(props(Room="103", students=10))]
 ///     Mathematics,
-///     #[strum(props(Time="2:30"))]
+///     #[strum(props(Time="2:30", mandatory=true))]
 ///     Science,
 /// }
 ///
 /// let history = Class::History;
 /// assert_eq!("Ms.Frizzle", history.get_str("Teacher").unwrap());
+/// assert_eq!(16, history.get_int("students").unwrap());
+/// assert!(history.get_bool("mandatory").unwrap());
 /// ```
 pub trait EnumProperty {
     fn get_str(&self, prop: &str) -> Option<&'static str>;
-    fn get_int(&self, _prop: &str) -> Option<usize> {
-        Option::None
-    }
-
-    fn get_bool(&self, _prop: &str) -> Option<bool> {
-        Option::None
-    }
+    fn get_int(&self, _prop: &str) -> Option<i64>;
+    fn get_bool(&self, _prop: &str) -> Option<bool>;
 }
 
 /// A cheap reference-to-reference conversion. Used to convert a value to a
diff --git a/strum_macros/src/helpers/metadata.rs b/strum_macros/src/helpers/metadata.rs
index 40d6394..4c333e4 100644
--- a/strum_macros/src/helpers/metadata.rs
+++ b/strum_macros/src/helpers/metadata.rs
@@ -204,7 +204,7 @@ pub enum VariantMeta {
     },
     Props {
         kw: kw::props,
-        props: Vec<(LitStr, LitStr)>,
+        props: Vec<(LitStr, Lit)>,
     },
 }
 
@@ -267,7 +267,7 @@ impl Parse for VariantMeta {
     }
 }
 
-struct Prop(Ident, LitStr);
+struct Prop(Ident, Lit);
 
 impl Parse for Prop {
     fn parse(input: ParseStream) -> syn::Result<Self> {
diff --git a/strum_macros/src/helpers/mod.rs b/strum_macros/src/helpers/mod.rs
index a00e8e4..5d37ea6 100644
--- a/strum_macros/src/helpers/mod.rs
+++ b/strum_macros/src/helpers/mod.rs
@@ -16,7 +16,7 @@ use syn::spanned::Spanned;
 pub fn missing_parse_err_attr_error() -> syn::Error {
     syn::Error::new(
         Span::call_site(),
-        "`parse_err_ty` and `parse_err_fn` attribute is both required.",
+        "`parse_err_ty` and `parse_err_fn` attributes are both required.",
     )
 }
 
diff --git a/strum_macros/src/helpers/variant_props.rs b/strum_macros/src/helpers/variant_props.rs
index e0fffdc..ff98817 100644
--- a/strum_macros/src/helpers/variant_props.rs
+++ b/strum_macros/src/helpers/variant_props.rs
@@ -1,5 +1,5 @@
 use std::default::Default;
-use syn::{Ident, LitStr, Variant};
+use syn::{Ident, Lit, LitStr, Variant};
 
 use super::case_style::{CaseStyle, CaseStyleHelpers};
 use super::metadata::{kw, VariantExt, VariantMeta};
@@ -18,7 +18,7 @@ pub struct StrumVariantProperties {
     pub message: Option<LitStr>,
     pub detailed_message: Option<LitStr>,
     pub documentation: Vec<LitStr>,
-    pub string_props: Vec<(LitStr, LitStr)>,
+    pub props: Vec<(LitStr, Lit)>,
     serialize: Vec<LitStr>,
     pub to_string: Option<LitStr>,
     ident: Option<Ident>,
@@ -143,7 +143,7 @@ impl HasStrumVariantProperties for Variant {
                     output.ascii_case_insensitive = Some(value);
                 }
                 VariantMeta::Props { props, .. } => {
-                    output.string_props.extend(props);
+                    output.props.extend(props);
                 }
             }
         }
diff --git a/strum_macros/src/macros/enum_properties.rs b/strum_macros/src/macros/enum_properties.rs
index b2a2b23..4ba9af3 100644
--- a/strum_macros/src/macros/enum_properties.rs
+++ b/strum_macros/src/macros/enum_properties.rs
@@ -1,9 +1,24 @@
+use std::collections::HashMap;
+
 use proc_macro2::TokenStream;
 use quote::quote;
-use syn::{Data, DeriveInput, Fields};
+use syn::{Data, DeriveInput, Fields, Lit};
 
 use crate::helpers::{non_enum_error, HasStrumVariantProperties, HasTypeProperties};
 
+#[derive(Hash, PartialEq, Eq)]
+enum PropertyType {
+    String,
+    Integer,
+    Bool,
+}
+
+const PROPERTY_TYPES: [PropertyType; 3] = [
+    PropertyType::String,
+    PropertyType::Integer,
+    PropertyType::Bool,
+];
+
 pub fn enum_properties_inner(ast: &DeriveInput) -> syn::Result<TokenStream> {
     let name = &ast.ident;
     let (impl_generics, ty_generics, where_clause) = ast.generics.split_for_impl();
@@ -14,11 +29,12 @@ pub fn enum_properties_inner(ast: &DeriveInput) -> syn::Result<TokenStream> {
     let type_properties = ast.get_type_properties()?;
     let strum_module_path = type_properties.crate_module_path();
 
-    let mut arms = Vec::new();
+    let mut built_arms: HashMap<_, _> = PROPERTY_TYPES.iter().map(|p| (p, Vec::new())).collect();
+
     for variant in variants {
         let ident = &variant.ident;
         let variant_properties = variant.get_variant_properties()?;
-        let mut string_arms = Vec::new();
+        let mut arms: HashMap<_, _> = PROPERTY_TYPES.iter().map(|p| (p, Vec::new())).collect();
         // But you can disable the messages.
         if variant_properties.disabled.is_some() {
             continue;
@@ -30,33 +46,69 @@ pub fn enum_properties_inner(ast: &DeriveInput) -> syn::Result<TokenStream> {
             Fields::Named(..) => quote! { {..} },
         };
 
-        for (key, value) in variant_properties.string_props {
-            string_arms.push(quote! { #key => ::core::option::Option::Some( #value )});
-        }
+        for (key, value) in variant_properties.props {
+            let property_type = match value {
+                Lit::Str(..) => PropertyType::String,
+                Lit::Bool(..) => PropertyType::Bool,
+                Lit::Int(..) => PropertyType::Integer,
+                _ => todo!("TODO"),
+            };
 
-        string_arms.push(quote! { _ => ::core::option::Option::None });
+            arms.get_mut(&property_type)
+                .unwrap()
+                .push(quote! { #key => ::core::option::Option::Some( #value )});
+        }
 
-        arms.push(quote! {
-            &#name::#ident #params => {
-                match prop {
-                    #(#string_arms),*
+        for property in &PROPERTY_TYPES {
+            arms.get_mut(&property)
+                .unwrap()
+                .push(quote! { _ => ::core::option::Option::None });
+            let arms_as_string = &arms[property];
+            built_arms.get_mut(&property).unwrap().push(quote! {
+                &#name::#ident #params => {
+                    match prop {
+                        #(#arms_as_string),*
+                    }
                 }
-            }
-        });
+            });
+        }
     }
 
-    if arms.len() < variants.len() {
-        arms.push(quote! { _ => ::core::option::Option::None });
+    for (_, arms) in built_arms.iter_mut() {
+        if arms.len() < variants.len() {
+            arms.push(quote! { _ => ::core::option::Option::None });
+        }
     }
 
+    let (built_string_arms, built_int_arms, built_bool_arms) = (
+        &built_arms[&PropertyType::String],
+        &built_arms[&PropertyType::Integer],
+        &built_arms[&PropertyType::Bool],
+    );
+
     Ok(quote! {
         impl #impl_generics #strum_module_path::EnumProperty for #name #ty_generics #where_clause {
             #[inline]
             fn get_str(&self, prop: &str) -> ::core::option::Option<&'static str> {
                 match self {
-                    #(#arms),*
+                    #(#built_string_arms),*
                 }
             }
+
+            #[inline]
+            fn get_int(&self, prop: &str) -> ::core::option::Option<i64> {
+                match self {
+                    #(#built_int_arms),*
+                }
+            }
+
+            #[inline]
+            fn get_bool(&self, prop: &str) -> ::core::option::Option<bool> {
+                match self {
+                    #(#built_bool_arms),*
+                }
+            }
+
         }
     })
 }


Test patch:
diff --git a/strum_tests/tests/enum_props.rs b/strum_tests/tests/enum_props.rs
index c32c07e..3a2ed32 100644
--- a/strum_tests/tests/enum_props.rs
+++ b/strum_tests/tests/enum_props.rs
@@ -47,3 +47,30 @@ fn crate_module_path_test() {
     let a = Test::A;
     assert_eq!("value", a.get_str("key").unwrap());
 }
+
+#[derive(Debug, EnumProperty)]
+enum TestGet {
+    #[strum(props(weight = 42, flat = true, big = false))]
+    A,
+    #[strum(props(weight = -42, flat = false))]
+    B,
+    C,
+}
+
+#[test]
+fn get_bool_test() {
+    const BOOL_KEY: &str = "flat";
+    assert_eq!(Some(true), TestGet::A.get_bool(BOOL_KEY));
+    assert_eq!(Some(false), TestGet::B.get_bool(BOOL_KEY));
+    assert_eq!(None, TestGet::C.get_bool(BOOL_KEY));
+    assert_eq!(None, TestGet::A.get_bool("weight"));
+}
+
+#[test]
+fn get_int_test() {
+    const INT_KEY: &str = "weight";
+    assert_eq!(Some(42), TestGet::A.get_int(INT_KEY));
+    assert_eq!(Some(-42), TestGet::B.get_int(INT_KEY));
+    assert_eq!(None, TestGet::C.get_int(INT_KEY));
+    assert_eq!(None, TestGet::A.get_int("flat"));
+}