Repository: Peternator7/strum
Pull request #74 (2019-11-02 18:41:39)

Move EnumVariantsNames into a trait which the macro implements
I ran into a case where I had to accept a generic parameter of any enum type that has the `variants()` function implemented, which would've previously required some awkward workarounds. This PR moves the function into its own trait `VariantNames` and changes the `EnumVariantNames` to implement it, so that it is possible to use the trait to accept any such enum (see the added test case).


Patch:
diff --git a/strum/src/lib.rs b/strum/src/lib.rs
index 67f33ab..fb36882 100644
--- a/strum/src/lib.rs
+++ b/strum/src/lib.rs
@@ -40,7 +40,7 @@
 //! | [Display] | Converts enum variants to strings |
 //! | [AsRefStr] | Converts enum variants to `&'static str` |
 //! | [IntoStaticStr] | Implements `From<MyEnum> for &'static str` on an enum |
-//! | [EnumVariantNames] | Adds a `variants` method returning an array of discriminant names |
+//! | [EnumVariantNames] | Implements Strum::VariantNames which adds a static `variants` method returning an array of discriminant names |
 //! | [EnumIter] | Creates a new type that iterates of the variants of an enum. |
 //! | [EnumProperty] | Add custom properties to enum variants. |
 //! | [EnumMessage] | Add a verbose message to an enum variant. |
@@ -218,6 +218,12 @@ pub trait EnumCount {
     fn count() -> usize;
 }
 
+/// A trait for retrieving the names of each variant in Enum. This trait can
+/// be autoderived by `strum_macros`.
+pub trait VariantNames {
+    fn variants() -> &'static [&'static str];
+}
+
 #[cfg(feature = "derive")]
 #[allow(unused_imports)]
 #[macro_use]
diff --git a/strum_macros/src/macros/enum_iter.rs b/strum_macros/src/macros/enum_iter.rs
index afc6d8b..a62bfd8 100644
--- a/strum_macros/src/macros/enum_iter.rs
+++ b/strum_macros/src/macros/enum_iter.rs
@@ -83,12 +83,14 @@ pub fn enum_iter_inner(ast: &syn::DeriveInput) -> TokenStream {
                     #(#arms),*
                 };
 
-                self.idx += 1;
+                if self.idx < #variant_count {
+                    self.idx += 1;
+                }
                 output
             }
 
             fn size_hint(&self) -> (usize, Option<usize>) {
-                let t = #variant_count - self.idx;
+                let t = if self.idx >= #variant_count { 0 } else { #variant_count - self.idx };
                 (t, Some(t))
             }
         }
diff --git a/strum_macros/src/macros/enum_variant_names.rs b/strum_macros/src/macros/enum_variant_names.rs
index a1d436d..0bc550e 100644
--- a/strum_macros/src/macros/enum_variant_names.rs
+++ b/strum_macros/src/macros/enum_variant_names.rs
@@ -1,11 +1,12 @@
 use proc_macro2::TokenStream;
 use syn;
 
-use crate::helpers::case_style::CaseStyle;
-use crate::helpers::{extract_meta, CaseStyleHelpers, MetaIteratorHelpers};
+use crate::helpers::{case_style::CaseStyle, extract_meta, CaseStyleHelpers, MetaIteratorHelpers};
 
 pub fn enum_variant_names_inner(ast: &syn::DeriveInput) -> TokenStream {
     let name = &ast.ident;
+    let gen = &ast.generics;
+    let (impl_generics, ty_generics, where_clause) = gen.split_for_impl();
 
     let variants = match ast.data {
         syn::Data::Enum(ref v) => &v.variants,
@@ -24,10 +25,10 @@ pub fn enum_variant_names_inner(ast: &syn::DeriveInput) -> TokenStream {
         .collect::<Vec<_>>();
 
     quote! {
-        impl #name {
+        impl #impl_generics ::strum::VariantNames for #name #ty_generics #where_clause {
             /// Return a slice containing the names of the variants of this enum
             #[allow(dead_code)]
-            pub fn variants() -> &'static [&'static str] {
+            fn variants() -> &'static [&'static str] {
                 &[
                     #(#names),*
                 ]


Test patch:
diff --git a/strum_tests/tests/enum_iter.rs b/strum_tests/tests/enum_iter.rs
index e3cffa5..b0813ef 100644
--- a/strum_tests/tests/enum_iter.rs
+++ b/strum_tests/tests/enum_iter.rs
@@ -63,6 +63,9 @@ fn len_test() {
     i.next();
 
     assert_eq!(0, i.len());
+    i.next();
+
+    assert_eq!(0, i.size_hint().1.unwrap());
 }
 
 #[test]
diff --git a/strum_tests/tests/enum_variant_names.rs b/strum_tests/tests/enum_variant_names.rs
index 957c14b..f5f78ce 100644
--- a/strum_tests/tests/enum_variant_names.rs
+++ b/strum_tests/tests/enum_variant_names.rs
@@ -3,6 +3,7 @@ extern crate strum_macros;
 #[macro_use]
 extern crate structopt;
 extern crate strum;
+use strum::VariantNames;
 
 #[test]
 fn simple() {
@@ -17,6 +18,23 @@ fn simple() {
     assert_eq!(&Color::variants(), &["Red", "Blue", "Yellow"]);
 }
 
+#[test]
+fn variant_names_trait() {
+    #[allow(dead_code)]
+    #[derive(EnumVariantNames)]
+    enum Color {
+        Red,
+        Blue,
+        Yellow,
+    }
+
+    fn generic_function<T: VariantNames>() {
+        assert_eq!(T::variants(), &["Red", "Blue", "Yellow"]);
+    }
+
+    generic_function::<Color>();
+}
+
 #[test]
 fn plain_kebab() {
     #[allow(dead_code)]