Skip to content

Commit

Permalink
Enum extra attributes (#1966)
Browse files Browse the repository at this point in the history 
* 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 <[email protected]>

* Fixup

---------

Co-authored-by: Zander Milroy <[email protected]>
  • Loading branch information
@tyt2y3 @onefifth
tyt2y3 and onefifth authored Nov 8, 2023
1 parent 0dbfb42 commit 8ab24cf
Show file tree
Hide file tree
Showing 4 changed files with 102 additions and 7 deletions.
7 changes: 7 additions & 0 deletions sea-orm-cli/src/cli.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
2 changes: 2 additions & 0 deletions sea-orm-cli/src/commands/generate.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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);
Expand Down
84 changes: 80 additions & 4 deletions sea-orm-codegen/src/entity/active_enum.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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());
Expand All @@ -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)]
Expand All @@ -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};

Expand All @@ -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)]
Expand Down Expand Up @@ -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(),
Expand All @@ -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()
)
}
}
16 changes: 13 additions & 3 deletions sea-orm-codegen/src/entity/writer.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}

Expand Down Expand Up @@ -95,8 +96,8 @@ where
)
}

/// convert attributes argument to token stream
fn bonus_attributes<T, I>(attributes: I) -> TokenStream
/// convert *_extra_attributes argument to token stream
pub(crate) fn bonus_attributes<T, I>(attributes: I) -> TokenStream
where
T: Into<String>,
I: IntoIterator<Item = T>,
Expand Down Expand Up @@ -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 {
Expand All @@ -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,
}
}
Expand All @@ -175,6 +178,7 @@ impl EntityWriter {
&context.with_serde,
context.with_copy_enums,
&context.enum_extra_derives,
&context.enum_extra_attributes,
));
}
WriterOutput { files }
Expand Down Expand Up @@ -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);
Expand All @@ -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);
Expand Down

0 comments on commit 8ab24cf

Please sign in to comment.