From e0f81e9caa6cfc26d7fd585da0ba993c32ef70dd Mon Sep 17 00:00:00 2001 From: MendyBerger <12537668+MendyBerger@users.noreply.github.com> Date: Sun, 16 Jul 2023 22:59:38 -0400 Subject: [PATCH] Added EnumTryAs --- strum_macros/src/helpers/mod.rs | 2 + strum_macros/src/helpers/snakify.rs | 18 ++++++ strum_macros/src/lib.rs | 35 +++++++++++ strum_macros/src/macros/enum_is.rs | 20 +------ strum_macros/src/macros/enum_try_as.rs | 80 ++++++++++++++++++++++++++ strum_macros/src/macros/mod.rs | 1 + strum_tests/tests/enum_try_as.rs | 44 ++++++++++++++ 7 files changed, 181 insertions(+), 19 deletions(-) create mode 100644 strum_macros/src/helpers/snakify.rs create mode 100644 strum_macros/src/macros/enum_try_as.rs create mode 100644 strum_tests/tests/enum_try_as.rs diff --git a/strum_macros/src/helpers/mod.rs b/strum_macros/src/helpers/mod.rs index 11aebc8..69eeabe 100644 --- a/strum_macros/src/helpers/mod.rs +++ b/strum_macros/src/helpers/mod.rs @@ -1,9 +1,11 @@ pub use self::case_style::CaseStyleHelpers; +pub use self::snakify::snakify; pub use self::type_props::HasTypeProperties; pub use self::variant_props::HasStrumVariantProperties; pub mod case_style; mod metadata; +mod snakify; pub mod type_props; pub mod variant_props; diff --git a/strum_macros/src/helpers/snakify.rs b/strum_macros/src/helpers/snakify.rs new file mode 100644 index 0000000..9655544 --- /dev/null +++ b/strum_macros/src/helpers/snakify.rs @@ -0,0 +1,18 @@ +use heck::ToSnakeCase; + +/// 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 = 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() +} diff --git a/strum_macros/src/lib.rs b/strum_macros/src/lib.rs index a37ccdc..f09d135 100644 --- a/strum_macros/src/lib.rs +++ b/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 unnamed variants. +/// E.g. `Message.try_as_write()`. +/// +/// These methods will only be generated for unnamed 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 diff --git a/strum_macros/src/macros/enum_is.rs b/strum_macros/src/macros/enum_is.rs index bde3851..b69fa0c 100644 --- a/strum_macros/src/macros/enum_is.rs +++ b/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 { } .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 = 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() -} diff --git a/strum_macros/src/macros/enum_try_as.rs b/strum_macros/src/macros/enum_try_as.rs new file mode 100644 index 0000000..dac35b3 --- /dev/null +++ b/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 { + 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) -> Option<(#(#types),*)> { + match self { + #enum_name::#variant_name (#(#field_names),*) => Some((#(#field_names),*)), + _ => None + } + } + + #[must_use] + #[inline] + pub const fn #ref_fn_name(&self) -> Option<(#(&#types),*)> { + match self { + #enum_name::#variant_name (#(#field_names),*) => Some((#(#field_names),*)), + _ => None + } + } + + #[must_use] + #[inline] + pub fn #mut_fn_name(&mut self) -> 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()) +} diff --git a/strum_macros/src/macros/mod.rs b/strum_macros/src/macros/mod.rs index a44be08..8df8cd6 100644 --- a/strum_macros/src/macros/mod.rs +++ b/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; diff --git a/strum_tests/tests/enum_try_as.rs b/strum_tests/tests/enum_try_as.rs new file mode 100644 index 0000000..d9a8bb9 --- /dev/null +++ b/strum_tests/tests/enum_try_as.rs @@ -0,0 +1,44 @@ +use strum::EnumTryAs; + +#[derive(EnumTryAs)] +enum Foo { + Unnamed0(), + Unnamed1(u128), + Unnamed2(bool, String), + #[strum(disabled)] + #[allow(dead_code)] + Disabled(u32), +} + +#[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()); +}