109 lines
3.5 KiB
Rust

use heck::CamelCase;
use proc_macro2::TokenStream;
use quote::{format_ident, quote};
use sea_query::DynIden;
use crate::WithSerde;
#[derive(Clone, Debug)]
pub struct ActiveEnum {
pub(crate) enum_name: DynIden,
pub(crate) values: Vec<DynIden>,
}
impl ActiveEnum {
pub fn impl_active_enum(&self, with_serde: &WithSerde, with_copy_enums: bool) -> TokenStream {
let enum_name = &self.enum_name.to_string();
let enum_iden = format_ident!("{}", enum_name.to_camel_case());
let values: Vec<String> = self.values.iter().map(|v| v.to_string()).collect();
let variants = values.iter().map(|v| v.trim()).map(|v| {
if v.chars().next().map(char::is_numeric).unwrap_or(false) {
format_ident!("_{}", v)
} else {
format_ident!("{}", v.to_camel_case())
}
});
let extra_derive = with_serde.extra_derive();
let copy_derive = if with_copy_enums {
quote! { , Copy }
} else {
quote! {}
};
quote! {
#[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum #copy_derive #extra_derive)]
#[sea_orm(rs_type = "String", db_type = "Enum", enum_name = #enum_name)]
pub enum #enum_iden {
#(
#[sea_orm(string_value = #values)]
#variants,
)*
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use pretty_assertions::assert_eq;
use sea_query::{Alias, IntoIden};
#[test]
fn test_enum_variant_starts_with_number() {
assert_eq!(
ActiveEnum {
enum_name: Alias::new("media_type").into_iden(),
values: vec![
"UNKNOWN",
"BITMAP",
"DRAWING",
"AUDIO",
"VIDEO",
"MULTIMEDIA",
"OFFICE",
"TEXT",
"EXECUTABLE",
"ARCHIVE",
"3D",
]
.into_iter()
.map(|variant| Alias::new(variant).into_iden())
.collect(),
}
.impl_active_enum(&WithSerde::None, true)
.to_string(),
quote!(
#[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum, Copy)]
#[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,
#[sea_orm(string_value = "DRAWING")]
Drawing,
#[sea_orm(string_value = "AUDIO")]
Audio,
#[sea_orm(string_value = "VIDEO")]
Video,
#[sea_orm(string_value = "MULTIMEDIA")]
Multimedia,
#[sea_orm(string_value = "OFFICE")]
Office,
#[sea_orm(string_value = "TEXT")]
Text,
#[sea_orm(string_value = "EXECUTABLE")]
Executable,
#[sea_orm(string_value = "ARCHIVE")]
Archive,
#[sea_orm(string_value = "3D")]
_3D,
}
)
.to_string()
)
}
}