DeriveDisplay macro for enum (#1726)

* WIP

* WIP

* WIP

* changed from storing ident to storing TokenStream internally

* Push for code sharing

* fmt

* changing test target to string

* updated Display not showing extra single quote, fixed clippy

* removed DeriveDisplay for test case not using display trait

* added test cases for display_value attribute

* Remove useless clone

* Add comment

* Remove useless clone

* Rename method

* Refactoring

---------

Co-authored-by: Billy Chan <ccw.billy.123@gmail.com>

sea-orm-macros/src/derives/active_enum.rs

@@ -101,6 +101,10 @@ impl ActiveEnum {
                     } else if meta.path.is_ident("num_value") {
                         is_int = true;
                         num_value = Some(meta.value()?.parse::<LitInt>()?);
+                    } else if meta.path.is_ident("display_value") {
+                        // This is a placeholder to prevent the `display_value` proc_macro attribute of `DeriveDisplay`
+                        // to be considered unknown attribute parameter
+                        meta.value()?.parse::<LitStr>()?;
                     } else {
                         return Err(meta.error(format!(
                             "Unknown attribute parameter found: {:?}",
@@ -379,14 +383,6 @@ impl ActiveEnum {
                 }
             }
 
-            #[automatically_derived]
-            impl std::fmt::Display for #ident {
-                fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-                    let v: sea_orm::sea_query::Value = <Self as sea_orm::ActiveEnum>::to_value(&self).into();
-                    write!(f, "{}", v)
-                }
-            }
-
             #impl_not_u8
         )
     }

sea-orm-macros/src/derives/active_enum_display.rs

@@ -0,0 +1,104 @@
+use proc_macro2::TokenStream;
+use quote::{quote, quote_spanned, ToTokens};
+use syn::{LitInt, LitStr};
+
+enum Error {
+    InputNotEnum,
+    Syn(syn::Error),
+}
+
+struct Display {
+    ident: syn::Ident,
+    variants: Vec<DisplayVariant>,
+}
+
+struct DisplayVariant {
+    ident: syn::Ident,
+    display_value: TokenStream,
+}
+
+impl Display {
+    fn new(input: syn::DeriveInput) -> Result<Self, Error> {
+        let ident = input.ident;
+
+        let variant_vec = match input.data {
+            syn::Data::Enum(syn::DataEnum { variants, .. }) => variants,
+            _ => return Err(Error::InputNotEnum),
+        };
+
+        let mut variants = Vec::new();
+        for variant in variant_vec {
+            let mut display_value = variant.ident.to_string().to_token_stream();
+            for attr in variant.attrs.iter() {
+                if !attr.path().is_ident("sea_orm") {
+                    continue;
+                }
+                attr.parse_nested_meta(|meta| {
+                    if meta.path.is_ident("string_value") {
+                        meta.value()?.parse::<LitStr>()?;
+                    } else if meta.path.is_ident("num_value") {
+                        meta.value()?.parse::<LitInt>()?;
+                    } else if meta.path.is_ident("display_value") {
+                        display_value = meta.value()?.parse::<LitStr>()?.to_token_stream();
+                    } else {
+                        return Err(meta.error(format!(
+                            "Unknown attribute parameter found: {:?}",
+                            meta.path.get_ident()
+                        )));
+                    }
+
+                    Ok(())
+                })
+                .map_err(Error::Syn)?;
+            }
+            variants.push(DisplayVariant {
+                ident: variant.ident,
+                display_value,
+            });
+        }
+        Ok(Display { ident, variants })
+    }
+
+    fn expand(&self) -> syn::Result<TokenStream> {
+        let expanded_impl_active_enum_display = self.impl_active_enum_display();
+
+        Ok(expanded_impl_active_enum_display)
+    }
+
+    fn impl_active_enum_display(&self) -> TokenStream {
+        let Self { ident, variants } = self;
+
+        let variant_idents: Vec<_> = variants
+            .iter()
+            .map(|variant| variant.ident.clone())
+            .collect();
+
+        let variant_display: Vec<_> = variants
+            .iter()
+            .map(|variant| variant.display_value.to_owned())
+            .collect();
+
+        quote!(
+            #[automatically_derived]
+            impl std::fmt::Display for #ident {
+                fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+                    write!(f, "{}", match self {
+                        #( Self::#variant_idents => #variant_display, )*
+                    })
+                }
+            }
+        )
+    }
+}
+
+pub fn expand_derive_active_enum_display(input: syn::DeriveInput) -> syn::Result<TokenStream> {
+    let ident_span = input.ident.span();
+
+    match Display::new(input) {
+        Ok(model) => model.expand(),
+        Err(Error::InputNotEnum) => Ok(quote_spanned! {
+            ident_span => compile_error!("you can only derive EnumDisplay on enums");
+        }),
+        Err(Error::Syn(e)) => Err(e),
+    }
+}

sea-orm-macros/src/derives/mod.rs

@@ -1,4 +1,5 @@
 mod active_enum;
+mod active_enum_display;
 mod active_model;
 mod active_model_behavior;
 mod attributes;
@@ -19,6 +20,7 @@ mod util;
 mod value_type;
 
 pub use active_enum::*;
+pub use active_enum_display::*;
 pub use active_model::*;
 pub use active_model_behavior::*;
 pub use column::*;

sea-orm-macros/src/lib.rs

@@ -842,3 +842,13 @@ pub fn derive_value_type(input: TokenStream) -> TokenStream {
         Err(e) => e.to_compile_error().into(),
     }
 }
+
+#[cfg(feature = "derive")]
+#[proc_macro_derive(DeriveDisplay, attributes(sea_orm))]
+pub fn derive_active_enum_display(input: TokenStream) -> TokenStream {
+    let input = parse_macro_input!(input as DeriveInput);
+    match derives::expand_derive_active_enum_display(input) {
+        Ok(ts) => ts.into(),
+        Err(e) => e.to_compile_error().into(),
+    }
+}

src/entity/active_enum.rs

@@ -17,7 +17,7 @@ use sea_query::{DynIden, Expr, Nullable, SimpleExpr, Value, ValueType};
 /// use sea_orm::entity::prelude::*;
 ///
 /// // Using the derive macro
-/// #[derive(Debug, PartialEq, EnumIter, DeriveActiveEnum)]
+/// #[derive(Debug, PartialEq, EnumIter, DeriveActiveEnum, DeriveDisplay)]
 /// #[sea_orm(
 ///     rs_type = "String",
 ///     db_type = "String(Some(1))",
@@ -85,7 +85,7 @@ use sea_query::{DynIden, Expr, Nullable, SimpleExpr, Value, ValueType};
 /// use sea_orm::entity::prelude::*;
 ///
 /// // Define the `Category` active enum
-/// #[derive(Debug, Clone, PartialEq, EnumIter, DeriveActiveEnum)]
+/// #[derive(Debug, Clone, PartialEq, EnumIter, DeriveActiveEnum, DeriveDisplay)]
 /// #[sea_orm(rs_type = "String", db_type = "String(Some(1))")]
 /// pub enum Category {
 ///     #[sea_orm(string_value = "B")]
@@ -216,7 +216,7 @@ mod tests {
             }
         }
 
-        #[derive(Debug, PartialEq, Eq, EnumIter, DeriveActiveEnum)]
+        #[derive(Debug, PartialEq, Eq, EnumIter, DeriveActiveEnum, DeriveDisplay)]
         #[sea_orm(
             rs_type = "String",
             db_type = "String(Some(1))",
@@ -268,15 +268,15 @@ mod tests {
         );
         assert_eq!(Category::values(), DeriveCategory::values());
 
-        assert_eq!(format!("{}", DeriveCategory::Big), "'B'");
+        assert_eq!(format!("{}", DeriveCategory::Big), "Big");
-        assert_eq!(format!("{}", DeriveCategory::Small), "'S'");
+        assert_eq!(format!("{}", DeriveCategory::Small), "Small");
     }
 
     #[test]
     fn active_enum_derive_signed_integers() {
         macro_rules! test_num_value_int {
             ($ident: ident, $rs_type: expr, $db_type: expr, $col_def: ident) => {
-                #[derive(Debug, PartialEq, Eq, EnumIter, DeriveActiveEnum)]
+                #[derive(Debug, PartialEq, Eq, EnumIter, DeriveActiveEnum, DeriveDisplay)]
                 #[sea_orm(rs_type = $rs_type, db_type = $db_type)]
                 pub enum $ident {
                     #[sea_orm(num_value = -10)]
@@ -293,7 +293,7 @@ mod tests {
 
         macro_rules! test_fallback_int {
             ($ident: ident, $fallback_type: ident, $rs_type: expr, $db_type: expr, $col_def: ident) => {
-                #[derive(Debug, PartialEq, Eq, EnumIter, DeriveActiveEnum)]
+                #[derive(Debug, PartialEq, Eq, EnumIter, DeriveActiveEnum, DeriveDisplay)]
                 #[sea_orm(rs_type = $rs_type, db_type = $db_type)]
                 #[repr(i32)]
                 pub enum $ident {
@@ -325,9 +325,9 @@ mod tests {
 
                 assert_eq!($ident::db_type(), ColumnType::$col_def.def());
 
-                assert_eq!(format!("{}", $ident::Big), "1");
+                assert_eq!(format!("{}", $ident::Big), "Big");
-                assert_eq!(format!("{}", $ident::Small), "0");
+                assert_eq!(format!("{}", $ident::Small), "Small");
-                assert_eq!(format!("{}", $ident::Negative), "-10");
+                assert_eq!(format!("{}", $ident::Negative), "Negative");
             };
         }
 
@@ -346,7 +346,7 @@ mod tests {
     fn active_enum_derive_unsigned_integers() {
         macro_rules! test_num_value_uint {
             ($ident: ident, $rs_type: expr, $db_type: expr, $col_def: ident) => {
-                #[derive(Debug, PartialEq, Eq, EnumIter, DeriveActiveEnum)]
+                #[derive(Debug, PartialEq, Eq, EnumIter, DeriveActiveEnum, DeriveDisplay)]
                 #[sea_orm(rs_type = $rs_type, db_type = $db_type)]
                 pub enum $ident {
                     #[sea_orm(num_value = 1)]
@@ -361,7 +361,7 @@ mod tests {
 
         macro_rules! test_fallback_uint {
             ($ident: ident, $fallback_type: ident, $rs_type: expr, $db_type: expr, $col_def: ident) => {
-                #[derive(Debug, PartialEq, Eq, EnumIter, DeriveActiveEnum)]
+                #[derive(Debug, PartialEq, Eq, EnumIter, DeriveActiveEnum, DeriveDisplay)]
                 #[sea_orm(rs_type = $rs_type, db_type = $db_type)]
                 #[repr($fallback_type)]
                 pub enum $ident {
@@ -390,8 +390,8 @@ mod tests {
 
                 assert_eq!($ident::db_type(), ColumnType::$col_def.def());
 
-                assert_eq!(format!("{}", $ident::Big), "1");
+                assert_eq!(format!("{}", $ident::Big), "Big");
-                assert_eq!(format!("{}", $ident::Small), "0");
+                assert_eq!(format!("{}", $ident::Small), "Small");
             };
         }
 
@@ -459,7 +459,7 @@ mod tests {
             assert_eq!(PopOSTypos::try_from_value(&val.to_owned()), Ok(variant));
         }
 
-        #[derive(Clone, Debug, PartialEq, EnumIter, DeriveActiveEnum)]
+        #[derive(Clone, Debug, PartialEq, EnumIter, DeriveActiveEnum, DeriveDisplay)]
         #[sea_orm(
             rs_type = "String",
             db_type = "String(None)",

src/entity/prelude.rs

@@ -11,8 +11,8 @@ pub use crate::{
 #[cfg(feature = "macros")]
 pub use crate::{
     DeriveActiveEnum, DeriveActiveModel, DeriveActiveModelBehavior, DeriveColumn,
-    DeriveCustomColumn, DeriveEntity, DeriveEntityModel, DeriveIntoActiveModel, DeriveModel,
+    DeriveCustomColumn, DeriveDisplay, DeriveEntity, DeriveEntityModel, DeriveIntoActiveModel,
-    DerivePrimaryKey, DeriveRelatedEntity, DeriveRelation, FromJsonQueryResult,
+    DeriveModel, DerivePrimaryKey, DeriveRelatedEntity, DeriveRelation, FromJsonQueryResult,
 };
 
 pub use async_trait;

src/lib.rs

@@ -349,7 +349,7 @@ pub use schema::*;
 #[cfg(feature = "macros")]
 pub use sea_orm_macros::{
     DeriveActiveEnum, DeriveActiveModel, DeriveActiveModelBehavior, DeriveColumn,
-    DeriveCustomColumn, DeriveEntity, DeriveEntityModel, DeriveIntoActiveModel,
+    DeriveCustomColumn, DeriveDisplay, DeriveEntity, DeriveEntityModel, DeriveIntoActiveModel,
     DeriveMigrationName, DeriveModel, DerivePartialModel, DerivePrimaryKey, DeriveRelatedEntity,
     DeriveRelation, FromJsonQueryResult, FromQueryResult,
 };

tests/active_enum_tests.rs

@@ -798,4 +798,12 @@ mod tests {
             )
         );
     }
+
+    #[test]
+    fn display_test() {
+        assert_eq!(format!("{}", Tea::BreakfastTea), "BreakfastTea");
+        assert_eq!(format!("{}", DisplayTea::BreakfastTea), "Breakfast");
+        assert_eq!(format!("{}", Tea::EverydayTea), "EverydayTea");
+        assert_eq!(format!("{}", DisplayTea::EverydayTea), "Everyday");
+    }
 }

tests/common/features/sea_orm_active_enums.rs

@@ -18,7 +18,7 @@ pub enum Color {
     White,
 }
 
-#[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum)]
+#[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum, DeriveDisplay)]
 #[sea_orm(rs_type = "String", db_type = "Enum", enum_name = "tea")]
 pub enum Tea {
     #[sea_orm(string_value = "EverydayTea")]
@@ -53,3 +53,12 @@ pub enum MediaType {
     #[sea_orm(string_value = "3D")]
     _3D,
 }
+
+#[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum, DeriveDisplay)]
+#[sea_orm(rs_type = "String", db_type = "Enum", enum_name = "tea")]
+pub enum DisplayTea {
+    #[sea_orm(string_value = "EverydayTea", display_value = "Everyday")]
+    EverydayTea,
+    #[sea_orm(string_value = "BreakfastTea", display_value = "Breakfast")]
+    BreakfastTea,
+}