From 0dbfb42bb7f0132657607fc5cbd4854d8279ba99 Mon Sep 17 00:00:00 2001 From: Lev Khoroshansky Date: Wed, 8 Nov 2023 20:04:24 +0100 Subject: [PATCH] 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 | 7 ++++ sea-orm-cli/src/commands/generate.rs | 2 + sea-orm-codegen/src/entity/active_enum.rs | 50 +++++++++++++++++++++-- sea-orm-codegen/src/entity/writer.rs | 30 +++++++++----- 4 files changed, 74 insertions(+), 15 deletions(-) diff --git a/sea-orm-cli/src/cli.rs b/sea-orm-cli/src/cli.rs index ff88d30a4..cc92208c6 100644 --- a/sea-orm-cli/src/cli.rs +++ b/sea-orm-cli/src/cli.rs @@ -280,6 +280,13 @@ pub enum GenerateSubcommands { )] model_extra_attributes: Vec, + #[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, + #[arg( long, default_value = "false", diff --git a/sea-orm-cli/src/commands/generate.rs b/sea-orm-cli/src/commands/generate.rs index 7b7a89562..8e8994fc0 100644 --- a/sea-orm-cli/src/commands/generate.rs +++ b/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); diff --git a/sea-orm-codegen/src/entity/active_enum.rs b/sea-orm-codegen/src/entity/active_enum.rs index 30c9c5df6..95fd0a03b 100644 --- a/sea-orm-codegen/src/entity/active_enum.rs +++ b/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 = 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, "e! {}) .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() + } + } } diff --git a/sea-orm-codegen/src/entity/writer.rs b/sea-orm-codegen/src/entity/writer.rs index 6600f4624..695a3de08 100644 --- a/sea-orm-codegen/src/entity/writer.rs +++ b/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 -fn bonus_derive(model_extra_derives: I) -> TokenStream +/// Converts *_extra_derives argument to token stream +pub(crate) fn bonus_derive(extra_derives: I) -> TokenStream where T: Into, I: IntoIterator, { - model_extra_derives - .into_iter() - .map(Into::::into) - .fold(TokenStream::default(), |acc, derive| { + extra_derives.into_iter().map(Into::::into).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, model_extra_attributes: Vec, + enum_extra_derives: Vec, 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( - self.write_sea_orm_active_enums(&context.with_serde, context.with_copy_enums), - ); + files.push(self.write_sea_orm_active_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 {