Enum extra attributes (#1966)

* Feature: Adds sea-orm-cli generate `--enum-extra-attributes` option (#1952)

* Adds `--enum-extra-derives`

* Adds test_enum_extra_attributes

---------

Co-authored-by: Chris Tsang <chris.2y3@outlook.com>

* Fixup

---------

Co-authored-by: Zander Milroy <zander.milroy+github@gmail.com>

sea-orm-cli/src/cli.rs

@@ -287,6 +287,13 @@ pub enum GenerateSubcommands {
         )]
         enum_extra_derives: Vec<String>,
 
+        #[arg(
+            long,
+            value_delimiter = ',',
+            help = r#"Add extra attributes to generated enums, no need for `#[]` (comma separated), e.g. `--enum-extra-attributes 'serde(rename_all = "camelCase")','ts(export)'`"#
+        )]
+        enum_extra_attributes: Vec<String>,
+
         #[arg(
             long,
             default_value = "false",

sea-orm-cli/src/commands/generate.rs

@@ -32,6 +32,7 @@ pub async fn run_generate_command(
             model_extra_derives,
             model_extra_attributes,
             enum_extra_derives,
+            enum_extra_attributes,
             seaography,
         } => {
             if verbose {
@@ -182,6 +183,7 @@ pub async fn run_generate_command(
                 model_extra_derives,
                 model_extra_attributes,
                 enum_extra_derives,
+                enum_extra_attributes,
                 seaography,
             );
             let output = EntityTransformer::transform(table_stmts)?.generate(&writer_context);

sea-orm-codegen/src/entity/active_enum.rs

@@ -17,6 +17,7 @@ impl ActiveEnum {
         with_serde: &WithSerde,
         with_copy_enums: bool,
         extra_derives: &TokenStream,
+        extra_attributes: &TokenStream,
     ) -> TokenStream {
         let enum_name = &self.enum_name.to_string();
         let enum_iden = format_ident!("{}", enum_name.to_upper_camel_case());
@@ -39,6 +40,7 @@ impl ActiveEnum {
         quote! {
             #[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum #copy_derive #serde_derive #extra_derives)]
             #[sea_orm(rs_type = "String", db_type = "Enum", enum_name = #enum_name)]
+            #extra_attributes
             pub enum #enum_iden {
                 #(
                     #[sea_orm(string_value = #values)]
@@ -51,9 +53,8 @@ impl ActiveEnum {
 
 #[cfg(test)]
 mod tests {
-    use crate::entity::writer::bonus_derive;
-
     use super::*;
+    use crate::entity::writer::{bonus_attributes, bonus_derive};
     use pretty_assertions::assert_eq;
     use sea_query::{Alias, IntoIden};
 
@@ -79,7 +80,12 @@ mod tests {
                 .map(|variant| Alias::new(variant).into_iden())
                 .collect(),
             }
-            .impl_active_enum(&WithSerde::None, true, &quote! {})
+            .impl_active_enum(
+                &WithSerde::None,
+                true,
+                &TokenStream::new(),
+                &TokenStream::new(),
+            )
             .to_string(),
             quote!(
                 #[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum, Copy)]
@@ -126,7 +132,8 @@ mod tests {
             .impl_active_enum(
                 &WithSerde::None,
                 true,
-                &bonus_derive(["specta::Type", "ts_rs::TS"])
+                &bonus_derive(["specta::Type", "ts_rs::TS"]),
+                &TokenStream::new(),
             )
             .to_string(),
             build_generated_enum(),
@@ -147,4 +154,73 @@ mod tests {
             .to_string()
         }
     }
+
+    #[test]
+    fn test_enum_extra_attributes() {
+        assert_eq!(
+            ActiveEnum {
+                enum_name: Alias::new("coinflip_result_type").into_iden(),
+                values: vec!["HEADS", "TAILS"]
+                    .into_iter()
+                    .map(|variant| Alias::new(variant).into_iden())
+                    .collect(),
+            }
+            .impl_active_enum(
+                &WithSerde::None,
+                true,
+                &TokenStream::new(),
+                &bonus_attributes([r#"serde(rename_all = "camelCase")"#])
+            )
+            .to_string(),
+            quote!(
+                #[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum, Copy)]
+                #[sea_orm(
+                    rs_type = "String",
+                    db_type = "Enum",
+                    enum_name = "coinflip_result_type"
+                )]
+                #[serde(rename_all = "camelCase")]
+                pub enum CoinflipResultType {
+                    #[sea_orm(string_value = "HEADS")]
+                    Heads,
+                    #[sea_orm(string_value = "TAILS")]
+                    Tails,
+                }
+            )
+            .to_string()
+        );
+        assert_eq!(
+            ActiveEnum {
+                enum_name: Alias::new("coinflip_result_type").into_iden(),
+                values: vec!["HEADS", "TAILS"]
+                    .into_iter()
+                    .map(|variant| Alias::new(variant).into_iden())
+                    .collect(),
+            }
+            .impl_active_enum(
+                &WithSerde::None,
+                true,
+                &TokenStream::new(),
+                &bonus_attributes([r#"serde(rename_all = "camelCase")"#, "ts(export)"])
+            )
+            .to_string(),
+            quote!(
+                #[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum, Copy)]
+                #[sea_orm(
+                    rs_type = "String",
+                    db_type = "Enum",
+                    enum_name = "coinflip_result_type"
+                )]
+                #[serde(rename_all = "camelCase")]
+                #[ts(export)]
+                pub enum CoinflipResultType {
+                    #[sea_orm(string_value = "HEADS")]
+                    Heads,
+                    #[sea_orm(string_value = "TAILS")]
+                    Tails,
+                }
+            )
+            .to_string()
+        )
+    }
 }

sea-orm-codegen/src/entity/writer.rs

@@ -48,6 +48,7 @@ pub struct EntityWriterContext {
     pub(crate) model_extra_derives: TokenStream,
     pub(crate) model_extra_attributes: TokenStream,
     pub(crate) enum_extra_derives: TokenStream,
+    pub(crate) enum_extra_attributes: TokenStream,
     pub(crate) seaography: bool,
 }
 
@@ -95,8 +96,8 @@ where
     )
 }
 
-/// convert attributes argument to token stream
+/// convert *_extra_attributes argument to token stream
-fn bonus_attributes<T, I>(attributes: I) -> TokenStream
+pub(crate) fn bonus_attributes<T, I>(attributes: I) -> TokenStream
 where
     T: Into<String>,
     I: IntoIterator<Item = T>,
@@ -145,6 +146,7 @@ impl EntityWriterContext {
         model_extra_derives: Vec<String>,
         model_extra_attributes: Vec<String>,
         enum_extra_derives: Vec<String>,
+        enum_extra_attributes: Vec<String>,
         seaography: bool,
     ) -> Self {
         Self {
@@ -159,6 +161,7 @@ impl EntityWriterContext {
             model_extra_derives: bonus_derive(model_extra_derives),
             model_extra_attributes: bonus_attributes(model_extra_attributes),
             enum_extra_derives: bonus_derive(enum_extra_derives),
+            enum_extra_attributes: bonus_attributes(enum_extra_attributes),
             seaography,
         }
     }
@@ -175,6 +178,7 @@ impl EntityWriter {
                 &context.with_serde,
                 context.with_copy_enums,
                 &context.enum_extra_derives,
+                &context.enum_extra_attributes,
             ));
         }
         WriterOutput { files }
@@ -289,6 +293,7 @@ impl EntityWriter {
         with_serde: &WithSerde,
         with_copy_enums: bool,
         extra_derives: &TokenStream,
+        extra_attributes: &TokenStream,
     ) -> OutputFile {
         let mut lines = Vec::new();
         Self::write_doc_comment(&mut lines);
@@ -298,7 +303,12 @@ impl EntityWriter {
             .enums
             .values()
             .map(|active_enum| {
-                active_enum.impl_active_enum(with_serde, with_copy_enums, extra_derives)
+                active_enum.impl_active_enum(
+                    with_serde,
+                    with_copy_enums,
+                    extra_derives,
+                    extra_attributes,
+                )
             })
             .collect();
         Self::write(&mut lines, code_blocks);