1use heck::ToUpperCamelCase;
2use proc_macro2::TokenStream;
3use quote::{format_ident, quote};
4use sea_query::DynIden;
5use std::fmt::Write;
6
7use crate::{EntityFormat, WithSerde};
8
9#[derive(Clone, Debug)]
10pub struct ActiveEnum {
11 pub(crate) enum_name: DynIden,
12 pub(crate) values: Vec<DynIden>,
13}
14
15impl ActiveEnum {
16 pub fn impl_active_enum(
17 &self,
18 with_serde: &WithSerde,
19 with_copy_enums: bool,
20 extra_derives: &TokenStream,
21 extra_attributes: &TokenStream,
22 entity_format: EntityFormat,
23 ) -> TokenStream {
24 let enum_name = &self.enum_name.to_string();
25 let enum_iden = format_ident!("{}", enum_name.to_upper_camel_case());
26 let values: Vec<String> = self.values.iter().map(|v| v.to_string()).collect();
27 let variants = values.iter().map(|v| v.trim()).map(|v| {
28 if v.chars().next().map(char::is_numeric).unwrap_or(false) {
29 format_ident!("_{}", v)
30 } else {
31 let variant_name = if v.is_empty() {
32 println!("Warning: item in the enumeration '{enum_name}' is an empty string, it will be converted to `__EmptyString`. You can modify it later as needed.");
33 "__EmptyString".to_string()
34 } else {
35 v.to_upper_camel_case()
36 };
37 if variant_name.is_empty() {
38 println!("Warning: item '{v}' in the enumeration '{enum_name}' cannot be converted into a valid Rust enum member name. It will be converted to its corresponding UTF-8 encoding. You can modify it later as needed.");
39 let mut ss = String::new();
40 for c in v.chars() {
41 if c.len_utf8() > 1 {
42 write!(&mut ss, "{c}").unwrap();
43 } else {
44 write!(&mut ss, "U{:04X}", c as u32).unwrap();
45 }
46 }
47 format_ident!("{}", ss)
48 } else {
49 format_ident!("{}", variant_name)
50 }
51 }
52 });
53
54 let serde_derive = with_serde.extra_derive();
55 let copy_derive = if with_copy_enums {
56 quote! { , Copy }
57 } else {
58 quote! {}
59 };
60
61 if entity_format == EntityFormat::Frontend {
62 quote! {
63 #[derive(Debug, Clone, PartialEq, Eq #copy_derive #serde_derive #extra_derives)]
64 #extra_attributes
65 pub enum #enum_iden {
66 #(
67 #variants,
68 )*
69 }
70 }
71 } else {
72 quote! {
73 #[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum #copy_derive #serde_derive #extra_derives)]
74 #[sea_orm(rs_type = "String", db_type = "Enum", enum_name = #enum_name)]
75 #extra_attributes
76 pub enum #enum_iden {
77 #(
78 #[sea_orm(string_value = #values)]
79 #variants,
80 )*
81 }
82 }
83 }
84 }
85}
86
87#[cfg(test)]
88mod tests {
89 use super::*;
90 use crate::entity::writer::{bonus_attributes, bonus_derive};
91 use pretty_assertions::assert_eq;
92 use sea_query::{Alias, IntoIden};
93
94 #[test]
95 fn test_enum_variant_starts_with_number() {
96 assert_eq!(
97 ActiveEnum {
98 enum_name: Alias::new("media_type").into_iden(),
99 values: vec![
100 "UNKNOWN",
101 "BITMAP",
102 "DRAWING",
103 "AUDIO",
104 "VIDEO",
105 "MULTIMEDIA",
106 "OFFICE",
107 "TEXT",
108 "EXECUTABLE",
109 "ARCHIVE",
110 "3D",
111 ]
112 .into_iter()
113 .map(|variant| Alias::new(variant).into_iden())
114 .collect(),
115 }
116 .impl_active_enum(
117 &WithSerde::None,
118 true,
119 &TokenStream::new(),
120 &TokenStream::new(),
121 EntityFormat::Compact,
122 )
123 .to_string(),
124 quote!(
125 #[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum, Copy)]
126 #[sea_orm(rs_type = "String", db_type = "Enum", enum_name = "media_type")]
127 pub enum MediaType {
128 #[sea_orm(string_value = "UNKNOWN")]
129 Unknown,
130 #[sea_orm(string_value = "BITMAP")]
131 Bitmap,
132 #[sea_orm(string_value = "DRAWING")]
133 Drawing,
134 #[sea_orm(string_value = "AUDIO")]
135 Audio,
136 #[sea_orm(string_value = "VIDEO")]
137 Video,
138 #[sea_orm(string_value = "MULTIMEDIA")]
139 Multimedia,
140 #[sea_orm(string_value = "OFFICE")]
141 Office,
142 #[sea_orm(string_value = "TEXT")]
143 Text,
144 #[sea_orm(string_value = "EXECUTABLE")]
145 Executable,
146 #[sea_orm(string_value = "ARCHIVE")]
147 Archive,
148 #[sea_orm(string_value = "3D")]
149 _3D,
150 }
151 )
152 .to_string()
153 )
154 }
155
156 #[test]
157 fn test_enum_extra_derives() {
158 assert_eq!(
159 ActiveEnum {
160 enum_name: Alias::new("media_type").into_iden(),
161 values: vec!["UNKNOWN", "BITMAP",]
162 .into_iter()
163 .map(|variant| Alias::new(variant).into_iden())
164 .collect(),
165 }
166 .impl_active_enum(
167 &WithSerde::None,
168 true,
169 &bonus_derive(["specta::Type", "ts_rs::TS"]),
170 &TokenStream::new(),
171 EntityFormat::Compact,
172 )
173 .to_string(),
174 build_generated_enum(),
175 );
176
177 #[rustfmt::skip]
178 fn build_generated_enum() -> String {
179 quote!(
180 #[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum, Copy, specta :: Type, ts_rs :: TS)]
181 #[sea_orm(rs_type = "String", db_type = "Enum", enum_name = "media_type")]
182 pub enum MediaType {
183 #[sea_orm(string_value = "UNKNOWN")]
184 Unknown,
185 #[sea_orm(string_value = "BITMAP")]
186 Bitmap,
187 }
188 )
189 .to_string()
190 }
191 }
192
193 #[test]
194 fn test_enum_extra_attributes() {
195 assert_eq!(
196 ActiveEnum {
197 enum_name: Alias::new("coinflip_result_type").into_iden(),
198 values: vec!["HEADS", "TAILS"]
199 .into_iter()
200 .map(|variant| Alias::new(variant).into_iden())
201 .collect(),
202 }
203 .impl_active_enum(
204 &WithSerde::None,
205 true,
206 &TokenStream::new(),
207 &bonus_attributes([r#"serde(rename_all = "camelCase")"#]),
208 EntityFormat::Compact,
209 )
210 .to_string(),
211 quote!(
212 #[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum, Copy)]
213 #[sea_orm(
214 rs_type = "String",
215 db_type = "Enum",
216 enum_name = "coinflip_result_type"
217 )]
218 #[serde(rename_all = "camelCase")]
219 pub enum CoinflipResultType {
220 #[sea_orm(string_value = "HEADS")]
221 Heads,
222 #[sea_orm(string_value = "TAILS")]
223 Tails,
224 }
225 )
226 .to_string()
227 );
228 assert_eq!(
229 ActiveEnum {
230 enum_name: Alias::new("coinflip_result_type").into_iden(),
231 values: vec!["HEADS", "TAILS"]
232 .into_iter()
233 .map(|variant| Alias::new(variant).into_iden())
234 .collect(),
235 }
236 .impl_active_enum(
237 &WithSerde::None,
238 true,
239 &TokenStream::new(),
240 &bonus_attributes([r#"serde(rename_all = "camelCase")"#, "ts(export)"]),
241 EntityFormat::Compact,
242 )
243 .to_string(),
244 quote!(
245 #[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum, Copy)]
246 #[sea_orm(
247 rs_type = "String",
248 db_type = "Enum",
249 enum_name = "coinflip_result_type"
250 )]
251 #[serde(rename_all = "camelCase")]
252 #[ts(export)]
253 pub enum CoinflipResultType {
254 #[sea_orm(string_value = "HEADS")]
255 Heads,
256 #[sea_orm(string_value = "TAILS")]
257 Tails,
258 }
259 )
260 .to_string()
261 )
262 }
263
264 #[test]
265 fn test_enum_variant_utf8_encode() {
266 assert_eq!(
267 ActiveEnum {
268 enum_name: Alias::new("ty").into_iden(),
269 values: vec![
270 "Question",
271 "QuestionsAdditional",
272 "Answer",
273 "Other",
274 "/",
275 "//",
276 "A-B-C",
277 "你好",
278 ]
279 .into_iter()
280 .map(|variant| Alias::new(variant).into_iden())
281 .collect(),
282 }
283 .impl_active_enum(
284 &WithSerde::None,
285 true,
286 &TokenStream::new(),
287 &TokenStream::new(),
288 EntityFormat::Compact,
289 )
290 .to_string(),
291 quote!(
292 #[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum, Copy)]
293 #[sea_orm(rs_type = "String", db_type = "Enum", enum_name = "ty")]
294 pub enum Ty {
295 #[sea_orm(string_value = "Question")]
296 Question,
297 #[sea_orm(string_value = "QuestionsAdditional")]
298 QuestionsAdditional,
299 #[sea_orm(string_value = "Answer")]
300 Answer,
301 #[sea_orm(string_value = "Other")]
302 Other,
303 #[sea_orm(string_value = "/")]
304 U002F,
305 #[sea_orm(string_value = "//")]
306 U002FU002F,
307 #[sea_orm(string_value = "A-B-C")]
308 ABC,
309 #[sea_orm(string_value = "你好")]
310 你好,
311 }
312 )
313 .to_string()
314 )
315 }
316}