Support --enum-extra-derives (#1934)

* feat: Support `--enum-extra-derives`

* test: Enum extra derives

* style: Don't format derives

* fix: Put `quote!` inside a function with `#[rustfmt::skip]`

sea-orm-cli/src/cli.rs

@@ -280,6 +280,13 @@ pub enum GenerateSubcommands {
         )]
         model_extra_attributes: Vec<String>,
 
+        #[arg(
+            long,
+            value_delimiter = ',',
+            help = "Add extra derive macros to generated enums (comma separated), e.g. `--enum-extra-derives 'ts_rs::Ts','CustomDerive'`"
+        )]
+        enum_extra_derives: Vec<String>,
+
         #[arg(
             long,
             default_value = "false",

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

@@ -31,6 +31,7 @@ pub async fn run_generate_command(
             lib,
             model_extra_derives,
             model_extra_attributes,
+            enum_extra_derives,
             seaography,
         } => {
             if verbose {
@@ -180,6 +181,7 @@ pub async fn run_generate_command(
                 serde_skip_hidden_column,
                 model_extra_derives,
                 model_extra_attributes,
+                enum_extra_derives,
                 seaography,
             );
             let output = EntityTransformer::transform(table_stmts)?.generate(&writer_context);

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

@@ -12,7 +12,12 @@ pub struct ActiveEnum {
 }
 
 impl ActiveEnum {
-    pub fn impl_active_enum(&self, with_serde: &WithSerde, with_copy_enums: bool) -> TokenStream {
+    pub fn impl_active_enum(
+        &self,
+        with_serde: &WithSerde,
+        with_copy_enums: bool,
+        extra_derives: &TokenStream,
+    ) -> TokenStream {
         let enum_name = &self.enum_name.to_string();
         let enum_iden = format_ident!("{}", enum_name.to_upper_camel_case());
         let values: Vec<String> = self.values.iter().map(|v| v.to_string()).collect();
@@ -24,7 +29,7 @@ impl ActiveEnum {
             }
         });
 
-        let extra_derive = with_serde.extra_derive();
+        let serde_derive = with_serde.extra_derive();
         let copy_derive = if with_copy_enums {
             quote! { , Copy }
         } else {
@@ -32,7 +37,7 @@ impl ActiveEnum {
         };
 
         quote! {
-            #[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum #copy_derive #extra_derive)]
+            #[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)]
             pub enum #enum_iden {
                 #(
@@ -46,6 +51,8 @@ impl ActiveEnum {
 
 #[cfg(test)]
 mod tests {
+    use crate::entity::writer::bonus_derive;
+
     use super::*;
     use pretty_assertions::assert_eq;
     use sea_query::{Alias, IntoIden};
@@ -72,7 +79,7 @@ mod tests {
                 .map(|variant| Alias::new(variant).into_iden())
                 .collect(),
             }
-            .impl_active_enum(&WithSerde::None, true)
+            .impl_active_enum(&WithSerde::None, true, &quote! {})
             .to_string(),
             quote!(
                 #[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum, Copy)]
@@ -105,4 +112,39 @@ mod tests {
             .to_string()
         )
     }
+
+    #[test]
+    fn test_enum_extra_derives() {
+        assert_eq!(
+            ActiveEnum {
+                enum_name: Alias::new("media_type").into_iden(),
+                values: vec!["UNKNOWN", "BITMAP",]
+                    .into_iter()
+                    .map(|variant| Alias::new(variant).into_iden())
+                    .collect(),
+            }
+            .impl_active_enum(
+                &WithSerde::None,
+                true,
+                &bonus_derive(["specta::Type", "ts_rs::TS"])
+            )
+            .to_string(),
+            build_generated_enum(),
+        );
+
+        #[rustfmt::skip]
+        fn build_generated_enum() -> String {
+            quote!(
+                #[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum, Copy, specta :: Type, ts_rs :: TS)]
+                #[sea_orm(rs_type = "String", db_type = "Enum", enum_name = "media_type")]
+                pub enum MediaType {
+                    #[sea_orm(string_value = "UNKNOWN")]
+                    Unknown,
+                    #[sea_orm(string_value = "BITMAP")]
+                    Bitmap,
+                }
+            )
+            .to_string()
+        }
+    }
 }

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

@@ -47,6 +47,7 @@ pub struct EntityWriterContext {
     pub(crate) serde_skip_deserializing_primary_key: bool,
     pub(crate) model_extra_derives: TokenStream,
     pub(crate) model_extra_attributes: TokenStream,
+    pub(crate) enum_extra_derives: TokenStream,
     pub(crate) seaography: bool,
 }
 
@@ -79,19 +80,19 @@ impl WithSerde {
     }
 }
 
-/// Converts model_extra_derives argument to token stream
+/// Converts *_extra_derives argument to token stream
-fn bonus_derive<T, I>(model_extra_derives: I) -> TokenStream
+pub(crate) fn bonus_derive<T, I>(extra_derives: I) -> TokenStream
 where
     T: Into<String>,
     I: IntoIterator<Item = T>,
 {
-    model_extra_derives
+    extra_derives.into_iter().map(Into::<String>::into).fold(
-        .into_iter()
+        TokenStream::default(),
-        .map(Into::<String>::into)
+        |acc, derive| {
-        .fold(TokenStream::default(), |acc, derive| {
             let tokens: TokenStream = derive.parse().unwrap();
             quote! { #acc, #tokens }
-        })
+        },
+    )
 }
 
 /// convert attributes argument to token stream
@@ -143,6 +144,7 @@ impl EntityWriterContext {
         serde_skip_hidden_column: bool,
         model_extra_derives: Vec<String>,
         model_extra_attributes: Vec<String>,
+        enum_extra_derives: Vec<String>,
         seaography: bool,
     ) -> Self {
         Self {
@@ -156,6 +158,7 @@ impl EntityWriterContext {
             serde_skip_hidden_column,
             model_extra_derives: bonus_derive(model_extra_derives),
             model_extra_attributes: bonus_attributes(model_extra_attributes),
+            enum_extra_derives: bonus_derive(enum_extra_derives),
             seaography,
         }
     }
@@ -168,9 +171,11 @@ impl EntityWriter {
         files.push(self.write_index_file(context.lib));
         files.push(self.write_prelude());
         if !self.enums.is_empty() {
-            files.push(
+            files.push(self.write_sea_orm_active_enums(
-                self.write_sea_orm_active_enums(&context.with_serde, context.with_copy_enums),
+                &context.with_serde,
-            );
+                context.with_copy_enums,
+                &context.enum_extra_derives,
+            ));
         }
         WriterOutput { files }
     }
@@ -283,6 +288,7 @@ impl EntityWriter {
         &self,
         with_serde: &WithSerde,
         with_copy_enums: bool,
+        extra_derives: &TokenStream,
     ) -> OutputFile {
         let mut lines = Vec::new();
         Self::write_doc_comment(&mut lines);
@@ -291,7 +297,9 @@ impl EntityWriter {
         let code_blocks = self
             .enums
             .values()
-            .map(|active_enum| active_enum.impl_active_enum(with_serde, with_copy_enums))
+            .map(|active_enum| {
+                active_enum.impl_active_enum(with_serde, with_copy_enums, extra_derives)
+            })
             .collect();
         Self::write(&mut lines, code_blocks);
         OutputFile {