From 60281a20735561ba8c040785fb5a35130e4acdce Mon Sep 17 00:00:00 2001 From: Zander Milroy Date: Wed, 8 Nov 2023 14:23:14 -0500 Subject: [PATCH 1/2] 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 --- sea-orm-cli/src/cli.rs | 7 +++ sea-orm-cli/src/commands/generate.rs | 2 + sea-orm-codegen/src/entity/active_enum.rs | 72 ++++++++++++++++++++++- sea-orm-codegen/src/entity/writer.rs | 11 +++- 4 files changed, 88 insertions(+), 4 deletions(-) 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..76e59ca22 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)] @@ -54,6 +56,7 @@ mod tests { use crate::entity::writer::bonus_derive; use super::*; + use crate::entity::writer::bonus_attributes; use pretty_assertions::assert_eq; use sea_query::{Alias, IntoIden}; @@ -79,7 +82,7 @@ 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)] @@ -147,4 +150,71 @@ 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, + &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, + &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..3daf5f118 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,7 @@ 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); From 06d10a7d36950541bde30945b2a6f4858c121875 Mon Sep 17 00:00:00 2001 From: Chris Tsang Date: Wed, 8 Nov 2023 19:27:17 +0000 Subject: [PATCH 2/2] Fixup --- sea-orm-codegen/src/entity/active_enum.rs | 16 +++++++++++----- sea-orm-codegen/src/entity/writer.rs | 7 ++++++- 2 files changed, 17 insertions(+), 6 deletions(-) diff --git a/sea-orm-codegen/src/entity/active_enum.rs b/sea-orm-codegen/src/entity/active_enum.rs index 76e59ca22..e7e0480c7 100644 --- a/sea-orm-codegen/src/entity/active_enum.rs +++ b/sea-orm-codegen/src/entity/active_enum.rs @@ -53,10 +53,8 @@ impl ActiveEnum { #[cfg(test)] mod tests { - use crate::entity::writer::bonus_derive; - use super::*; - use crate::entity::writer::bonus_attributes; + use crate::entity::writer::{bonus_attributes, bonus_derive}; use pretty_assertions::assert_eq; use sea_query::{Alias, IntoIden}; @@ -82,7 +80,12 @@ mod tests { .map(|variant| Alias::new(variant).into_iden()) .collect(), } - .impl_active_enum(&WithSerde::None, true, &TokenStream::new(), &TokenStream::new()) + .impl_active_enum( + &WithSerde::None, + true, + &TokenStream::new(), + &TokenStream::new(), + ) .to_string(), quote!( #[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum, Copy)] @@ -129,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(), @@ -164,6 +168,7 @@ mod tests { .impl_active_enum( &WithSerde::None, true, + &TokenStream::new(), &bonus_attributes([r#"serde(rename_all = "camelCase")"#]) ) .to_string(), @@ -195,6 +200,7 @@ mod tests { .impl_active_enum( &WithSerde::None, true, + &TokenStream::new(), &bonus_attributes([r#"serde(rename_all = "camelCase")"#, "ts(export)"]) ) .to_string(), diff --git a/sea-orm-codegen/src/entity/writer.rs b/sea-orm-codegen/src/entity/writer.rs index 3daf5f118..099a96797 100644 --- a/sea-orm-codegen/src/entity/writer.rs +++ b/sea-orm-codegen/src/entity/writer.rs @@ -303,7 +303,12 @@ impl EntityWriter { .enums .values() .map(|active_enum| { - active_enum.impl_active_enum(with_serde, with_copy_enums, extra_derives, extra_attributes) + active_enum.impl_active_enum( + with_serde, + with_copy_enums, + extra_derives, + extra_attributes, + ) }) .collect(); Self::write(&mut lines, code_blocks);