Enum extra attributes (#1966)

* 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 <chris.2y3@outlook.com>

* Fixup

---------

Co-authored-by: Zander Milroy <zander.milroy+github@gmail.com>
This commit is contained in:
Chris Tsang 2023-11-08 19:41:57 +00:00 committed by GitHub
parent 0dbfb42bb7
commit 8ab24cf884
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 102 additions and 7 deletions

View File

@ -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",

View File

@ -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);

View File

@ -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, &quote! {})
.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()
)
}
}

View File

@ -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<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>,
@ -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 {
@ -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);