diff --git a/sea-orm-cli/src/cli.rs b/sea-orm-cli/src/cli.rs index cc92208c6..e8611373e 100644 --- a/sea-orm-cli/src/cli.rs +++ b/sea-orm-cli/src/cli.rs @@ -287,6 +287,13 @@ pub enum GenerateSubcommands { )] enum_extra_derives: Vec, + #[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, + #[arg( long, default_value = "false", diff --git a/sea-orm-cli/src/commands/generate.rs b/sea-orm-cli/src/commands/generate.rs index 8e8994fc0..613a541fb 100644 --- a/sea-orm-cli/src/commands/generate.rs +++ b/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); diff --git a/sea-orm-codegen/src/entity/active_enum.rs b/sea-orm-codegen/src/entity/active_enum.rs index 95fd0a03b..e7e0480c7 100644 --- a/sea-orm-codegen/src/entity/active_enum.rs +++ b/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, "e! {}) + .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() + ) + } } diff --git a/sea-orm-codegen/src/entity/writer.rs b/sea-orm-codegen/src/entity/writer.rs index 695a3de08..099a96797 100644 --- a/sea-orm-codegen/src/entity/writer.rs +++ b/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 -fn bonus_attributes(attributes: I) -> TokenStream +/// convert *_extra_attributes argument to token stream +pub(crate) fn bonus_attributes(attributes: I) -> TokenStream where T: Into, I: IntoIterator, @@ -145,6 +146,7 @@ impl EntityWriterContext { model_extra_derives: Vec, model_extra_attributes: Vec, enum_extra_derives: Vec, + enum_extra_attributes: Vec, 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);