Merge branch 'master' of github.com:Peternator7/strum

README.md

@@ -380,6 +380,11 @@ Strum has implemented the following macros:
     }
     ```
 
+8. `EnumCount`: for a given enum generates implementation of `strum::EnumCount`,
+    which returns number of variants via `strum::EnumCount::count` method,
+	also for given `enum MyEnum` generates `const MYENUM_COUNT: usize`
+	which gives the same value as `strum::EnumCount` (which is usefull for array sizes, etc.).
+
 # Additional Attributes
 
 Strum supports several custom attributes to modify the generated code. At the enum level, the

strum/src/lib.rs

@@ -669,3 +669,8 @@ where
 {
     fn as_static(&self) -> &'static T;
 }
+
+/// Number of variants in Enum
+pub trait EnumCount {
+    fn count() -> usize;
+}

strum_macros/src/enum_count.rs

@@ -0,0 +1,30 @@
+use proc_macro2::{Span, TokenStream};
+use syn;
+
+pub(crate) fn enum_count_inner(ast: &syn::DeriveInput) -> TokenStream {
+    let n = match ast.data {
+        syn::Data::Enum(ref v) => v.variants.len(),
+        _ => panic!("EnumCount can only be used with enums"),
+    };
+    // Used in the quasi-quotation below as `#name`
+    let name = &ast.ident;
+    let const_name = &syn::Ident::new(
+        &format!("{}_COUNT", name.to_string().to_uppercase()),
+        Span::call_site(),
+    );
+
+    // Helper is provided for handling complex generic types correctly and effortlessly
+    let (impl_generics, ty_generics, where_clause) = ast.generics.split_for_impl();
+
+    quote! {
+            // The generated impl
+            impl #impl_generics ::strum::EnumCount for #name #ty_generics #where_clause {
+                fn count() -> usize {
+                    #n
+                }
+            }
+
+            #[allow(dead_code)]
+            pub const #const_name: usize = #n;
+    }
+}

strum_macros/src/lib.rs

@@ -20,6 +20,7 @@ extern crate proc_macro2;
 mod as_ref_str;
 mod case_style;
 mod display;
+mod enum_count;
 mod enum_discriminants;
 mod enum_iter;
 mod enum_messages;
@@ -124,3 +125,11 @@ pub fn enum_discriminants(input: proc_macro::TokenStream) -> proc_macro::TokenSt
     debug_print_generated(&ast, &toks);
     toks.into()
 }
+
+#[proc_macro_derive(EnumCount, attributes(strum))]
+pub fn enum_count(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
+    let ast = syn::parse(input).unwrap();
+    let toks = enum_count::enum_count_inner(&ast);
+    debug_print_generated(&ast, &toks);
+    toks.into()
+}

strum_tests/tests/enum_count.rs

@@ -0,0 +1,23 @@
+extern crate strum;
+#[macro_use]
+extern crate strum_macros;
+
+use strum::{EnumCount, IntoEnumIterator};
+
+#[derive(Debug, EnumCount, EnumIter)]
+enum Week {
+    Sunday,
+    Monday,
+    Tuesday,
+    Wednesday,
+    Thursday,
+    Friday,
+    Saturday,
+}
+
+#[test]
+fn simple_test() {
+    assert_eq!(7, Week::count());
+    assert_eq!(Week::count(), WEEK_COUNT);
+    assert_eq!(Week::iter().count(), WEEK_COUNT);
+}