Merge pull request Peternator7#285 from MendyBerger:try-as

Added EnumTryAs

strum_macros/src/helpers/case_style.rs

@@ -158,3 +158,20 @@ mod tests {
         assert_eq!(MixedCase, f("mixed_case").unwrap());
     }
 }
+
+/// heck doesn't treat numbers as new words, but this function does.
+/// E.g. for input `Hello2You`, heck would output `hello2_you`, and snakify would output `hello_2_you`.
+pub fn snakify(s: &str) -> String {
+    let mut output: Vec<char> = s.to_string().to_snake_case().chars().collect();
+    let mut num_starts = vec![];
+    for (pos, c) in output.iter().enumerate() {
+        if c.is_digit(10) && pos != 0 && !output[pos - 1].is_digit(10) {
+            num_starts.push(pos);
+        }
+    }
+    // need to do in reverse, because after inserting, all chars after the point of insertion are off
+    for i in num_starts.into_iter().rev() {
+        output.insert(i, '_')
+    }
+    output.into_iter().collect()
+}

strum_macros/src/helpers/mod.rs

@@ -1,4 +1,4 @@
-pub use self::case_style::CaseStyleHelpers;
+pub use self::case_style::{CaseStyleHelpers, snakify};
 pub use self::type_props::HasTypeProperties;
 pub use self::variant_props::HasStrumVariantProperties;
 

strum_macros/src/lib.rs

@@ -409,6 +409,41 @@ pub fn enum_is(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
     toks.into()
 }
 
+/// Generated `try_as_*()` methods for all tuple-style variants.
+/// E.g. `Message.try_as_write()`.
+///
+/// These methods will only be generated for tuple-style variants, not for named or unit variants.
+///
+/// ```
+/// use strum_macros::EnumTryAs;
+///
+/// #[derive(EnumTryAs, Debug)]
+/// enum Message {
+///     Quit,
+///     Move { x: i32, y: i32 },
+///     Write(String),
+///     ChangeColor(i32, i32, i32),
+/// }
+///
+/// assert_eq!(
+///     Message::Write(String::from("Hello")).try_as_write(),
+///     Some(String::from("Hello"))
+/// );
+/// assert_eq!(
+///     Message::ChangeColor(1, 2, 3).try_as_change_color(),
+///     Some((1, 2, 3))
+/// );
+/// ```
+#[proc_macro_derive(EnumTryAs, attributes(strum))]
+pub fn enum_try_as(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
+    let ast = syn::parse_macro_input!(input as DeriveInput);
+
+    let toks =
+        macros::enum_try_as::enum_try_as_inner(&ast).unwrap_or_else(|err| err.to_compile_error());
+    debug_print_generated(&ast, &toks);
+    toks.into()
+}
+
 /// Add a function to enum that allows accessing variants by its discriminant
 ///
 /// This macro adds a standalone function to obtain an enum variant by its discriminant. The macro adds

strum_macros/src/macros/enum_is.rs

@@ -1,5 +1,4 @@
-use crate::helpers::{non_enum_error, HasStrumVariantProperties};
-use heck::ToSnakeCase;
+use crate::helpers::{non_enum_error, snakify, HasStrumVariantProperties};
 use proc_macro2::TokenStream;
 use quote::{format_ident, quote};
 use syn::{Data, DeriveInput};
@@ -42,20 +41,3 @@ pub fn enum_is_inner(ast: &DeriveInput) -> syn::Result<TokenStream> {
     }
     .into())
 }
-
-/// heck doesn't treat numbers as new words, but this function does.
-/// E.g. for input `Hello2You`, heck would output `hello2_you`, and snakify would output `hello_2_you`.
-fn snakify(s: &str) -> String {
-    let mut output: Vec<char> = s.to_string().to_snake_case().chars().collect();
-    let mut num_starts = vec![];
-    for (pos, c) in output.iter().enumerate() {
-        if c.is_digit(10) && pos != 0 && !output[pos - 1].is_digit(10) {
-            num_starts.push(pos);
-        }
-    }
-    // need to do in reverse, because after inserting, all chars after the point of insertion are off
-    for i in num_starts.into_iter().rev() {
-        output.insert(i, '_')
-    }
-    output.into_iter().collect()
-}

strum_macros/src/macros/enum_try_as.rs

@@ -0,0 +1,80 @@
+use crate::helpers::{non_enum_error, snakify, HasStrumVariantProperties};
+use proc_macro2::TokenStream;
+use quote::{format_ident, quote, ToTokens};
+use syn::{Data, DeriveInput};
+
+pub fn enum_try_as_inner(ast: &DeriveInput) -> syn::Result<TokenStream> {
+    let variants = match &ast.data {
+        Data::Enum(v) => &v.variants,
+        _ => return Err(non_enum_error()),
+    };
+
+    let enum_name = &ast.ident;
+
+    let variants: Vec<_> = variants
+        .iter()
+        .filter_map(|variant| {
+            if variant.get_variant_properties().ok()?.disabled.is_some() {
+                return None;
+            }
+
+            match &variant.fields {
+                syn::Fields::Unnamed(values) => {
+                    let variant_name = &variant.ident;
+                    let types: Vec<_> = values.unnamed.iter().map(|field| {
+                        field.to_token_stream()
+                    }).collect();
+                    let field_names: Vec<_> = values.unnamed.iter().enumerate().map(|(i, _)| {
+                        let name = "x".repeat(i + 1);
+                        let name = format_ident!("{}", name);
+                        quote! {#name}
+                    }).collect();
+
+                    let move_fn_name = format_ident!("try_as_{}", snakify(&variant_name.to_string()));
+                    let ref_fn_name = format_ident!("try_as_{}_ref", snakify(&variant_name.to_string()));
+                    let mut_fn_name = format_ident!("try_as_{}_mut", snakify(&variant_name.to_string()));
+
+                    Some(quote! {
+                        #[must_use]
+                        #[inline]
+                        pub fn #move_fn_name(self) -> ::core::option::Option<(#(#types),*)> {
+                            match self {
+                                #enum_name::#variant_name (#(#field_names),*) => Some((#(#field_names),*)),
+                                _ => None
+                            }
+                        }
+
+                        #[must_use]
+                        #[inline]
+                        pub const fn #ref_fn_name(&self) -> ::core::option::Option<(#(&#types),*)> {
+                            match self {
+                                #enum_name::#variant_name (#(#field_names),*) => Some((#(#field_names),*)),
+                                _ => None
+                            }
+                        }
+
+                        #[must_use]
+                        #[inline]
+                        pub fn #mut_fn_name(&mut self) -> ::core::option::Option<(#(&mut #types),*)> {
+                            match self {
+                                #enum_name::#variant_name (#(#field_names),*) => Some((#(#field_names),*)),
+                                _ => None
+                            }
+                        }
+                    })
+                },
+                _ => {
+                    return None;
+                }
+            }
+
+        })
+        .collect();
+
+    Ok(quote! {
+        impl #enum_name {
+            #(#variants)*
+        }
+    }
+    .into())
+}

strum_macros/src/macros/mod.rs

@@ -4,6 +4,7 @@ pub mod enum_is;
 pub mod enum_iter;
 pub mod enum_messages;
 pub mod enum_properties;
+pub mod enum_try_as;
 pub mod enum_variant_names;
 pub mod from_repr;
 

strum_tests/tests/enum_try_as.rs

@@ -0,0 +1,48 @@
+use strum::EnumTryAs;
+
+#[derive(EnumTryAs)]
+enum Foo {
+    Unnamed0(),
+    Unnamed1(u128),
+    Unnamed2(bool, String),
+    #[strum(disabled)]
+    #[allow(dead_code)]
+    Disabled(u32),
+    #[allow(dead_code)]
+    Unit,
+    #[allow(dead_code)]
+    Named { _a: u32, _b: String },
+}
+
+#[test]
+fn unnamed_0() {
+    let foo = Foo::Unnamed0();
+    assert_eq!(Some(()), foo.try_as_unnamed_0());
+}
+
+#[test]
+fn unnamed_1() {
+    let foo = Foo::Unnamed1(128);
+    assert_eq!(Some(&128), foo.try_as_unnamed_1_ref());
+}
+
+#[test]
+fn unnamed_2() {
+    let foo = Foo::Unnamed2(true, String::from("Hay"));
+    assert_eq!(Some((true, String::from("Hay"))), foo.try_as_unnamed_2());
+}
+
+#[test]
+fn can_mutate() {
+    let mut foo = Foo::Unnamed1(128);
+    if let Some(value) = foo.try_as_unnamed_1_mut() {
+        *value = 44_u128;
+    }
+    assert_eq!(foo.try_as_unnamed_1(), Some(44));
+}
+
+#[test]
+fn doesnt_match_other_variations() {
+    let foo = Foo::Unnamed1(66);
+    assert_eq!(None, foo.try_as_unnamed_0());
+}