1use heck::ToUpperCamelCase;
2use proc_macro2::TokenStream;
3use quote::{format_ident, quote};
4use sea_query::DynIden;
5
6use crate::{EntityFormat, WithSerde};
7
8#[derive(Clone, Debug)]
9pub struct ActiveEnum {
10 pub(crate) enum_name: DynIden,
11 pub(crate) values: Vec<DynIden>,
12}
13
14impl ActiveEnum {
15 pub fn impl_active_enum(
16 &self,
17 with_serde: &WithSerde,
18 with_copy_enums: bool,
19 extra_derives: &TokenStream,
20 extra_attributes: &TokenStream,
21 entity_format: EntityFormat,
22 ) -> TokenStream {
23 let enum_name = &self.enum_name.to_string();
24 let enum_iden = format_ident!("{}", enum_name.to_upper_camel_case());
25 let values: Vec<String> = self.values.iter().map(|v| v.to_string()).collect();
26
27 let variants = values.iter().map(|v| v.trim()).map(|v| {
28 if v.is_empty() {
29 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.");
30 return format_ident!("__EmptyString");
31 }
32
33 let is_leading_digit = v.chars().next().is_some_and(|c| c.is_ascii_digit());
34 let is_valid_char = |c: char| c.is_ascii_alphanumeric() || c == '_';
35 let is_passthrough = |c: char| matches!(c, '-' | ' ');
36 let needs_utf8_escape = |c: char| !is_valid_char(c) && !is_passthrough(c);
37
38 if v.chars().any(needs_utf8_escape) {
39 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.");
40
41 let mut buf = String::new();
42
43 if is_leading_digit {
44 buf.push('_');
45 }
46
47 for c in v.chars() {
48 if is_passthrough(c) {
49 continue;
50 } else if is_valid_char(c) || c.len_utf8() > 1 {
51 buf.push(c);
52 } else {
53 buf.push_str(&format!("U{:04X}", c as u32));
54 }
55 }
56
57 return format_ident!("{buf}");
58 }
59
60 if is_leading_digit {
61 let sanitized = v.chars().filter(|&c| is_valid_char(c)).collect::<String>();
62 format_ident!("_{sanitized}")
63 } else {
64 format_ident!("{}", v.to_upper_camel_case())
65 }
66 });
67
68 let serde_derive = with_serde.extra_derive();
69 let copy_derive = if with_copy_enums {
70 quote! { , Copy }
71 } else {
72 quote! {}
73 };
74
75 if entity_format == EntityFormat::Frontend {
76 quote! {
77 #[derive(Debug, Clone, PartialEq, Eq #copy_derive #serde_derive #extra_derives)]
78 #extra_attributes
79 pub enum #enum_iden {
80 #(
81 #variants,
82 )*
83 }
84 }
85 } else {
86 quote! {
87 #[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum #copy_derive #serde_derive #extra_derives)]
88 #[sea_orm(rs_type = "String", db_type = "Enum", enum_name = #enum_name)]
89 #extra_attributes
90 pub enum #enum_iden {
91 #(
92 #[sea_orm(string_value = #values)]
93 #variants,
94 )*
95 }
96 }
97 }
98 }
99}
100
101#[cfg(test)]
102mod tests {
103 use super::*;
104 use crate::entity::writer::{bonus_attributes, bonus_derive};
105 use pretty_assertions::assert_eq;
106 use sea_query::{Alias, IntoIden};
107
108 #[test]
109 fn test_enum_variant_starts_with_empty_string() {
110 assert_eq!(
111 ActiveEnum {
112 enum_name: Alias::new("media_type").into_iden(),
113 values: vec![""]
114 .into_iter()
115 .map(|variant| Alias::new(variant).into_iden())
116 .collect(),
117 }
118 .impl_active_enum(
119 &WithSerde::None,
120 true,
121 &TokenStream::new(),
122 &TokenStream::new(),
123 EntityFormat::Compact,
124 )
125 .to_string(),
126 quote!(
127 #[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum, Copy)]
128 #[sea_orm(rs_type = "String", db_type = "Enum", enum_name = "media_type")]
129 pub enum MediaType {
130 #[sea_orm(string_value = "")]
131 __EmptyString,
132 }
133 )
134 .to_string()
135 )
136 }
137
138 #[test]
139 fn test_enum_variant_starts_with_number() {
140 assert_eq!(
141 ActiveEnum {
142 enum_name: Alias::new("media_type").into_iden(),
143 values: vec![
144 "UNKNOWN",
145 "BITMAP",
146 "DRAWING",
147 "AUDIO",
148 "VIDEO",
149 "MULTIMEDIA",
150 "OFFICE",
151 "TEXT",
152 "EXECUTABLE",
153 "ARCHIVE",
154 "3D",
155 "2-D"
156 ]
157 .into_iter()
158 .map(|variant| Alias::new(variant).into_iden())
159 .collect(),
160 }
161 .impl_active_enum(
162 &WithSerde::None,
163 true,
164 &TokenStream::new(),
165 &TokenStream::new(),
166 EntityFormat::Compact,
167 )
168 .to_string(),
169 quote!(
170 #[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum, Copy)]
171 #[sea_orm(rs_type = "String", db_type = "Enum", enum_name = "media_type")]
172 pub enum MediaType {
173 #[sea_orm(string_value = "UNKNOWN")]
174 Unknown,
175 #[sea_orm(string_value = "BITMAP")]
176 Bitmap,
177 #[sea_orm(string_value = "DRAWING")]
178 Drawing,
179 #[sea_orm(string_value = "AUDIO")]
180 Audio,
181 #[sea_orm(string_value = "VIDEO")]
182 Video,
183 #[sea_orm(string_value = "MULTIMEDIA")]
184 Multimedia,
185 #[sea_orm(string_value = "OFFICE")]
186 Office,
187 #[sea_orm(string_value = "TEXT")]
188 Text,
189 #[sea_orm(string_value = "EXECUTABLE")]
190 Executable,
191 #[sea_orm(string_value = "ARCHIVE")]
192 Archive,
193 #[sea_orm(string_value = "3D")]
194 _3D,
195 #[sea_orm(string_value = "2-D")]
196 _2D,
197 }
198 )
199 .to_string()
200 )
201 }
202
203 #[test]
204 fn test_enum_extra_derives() {
205 assert_eq!(
206 ActiveEnum {
207 enum_name: Alias::new("media_type").into_iden(),
208 values: vec!["UNKNOWN", "BITMAP",]
209 .into_iter()
210 .map(|variant| Alias::new(variant).into_iden())
211 .collect(),
212 }
213 .impl_active_enum(
214 &WithSerde::None,
215 true,
216 &bonus_derive(["specta::Type", "ts_rs::TS"]),
217 &TokenStream::new(),
218 EntityFormat::Compact,
219 )
220 .to_string(),
221 build_generated_enum(),
222 );
223
224 #[rustfmt::skip]
225 fn build_generated_enum() -> String {
226 quote!(
227 #[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum, Copy, specta :: Type, ts_rs :: TS)]
228 #[sea_orm(rs_type = "String", db_type = "Enum", enum_name = "media_type")]
229 pub enum MediaType {
230 #[sea_orm(string_value = "UNKNOWN")]
231 Unknown,
232 #[sea_orm(string_value = "BITMAP")]
233 Bitmap,
234 }
235 )
236 .to_string()
237 }
238 }
239
240 #[test]
241 fn test_enum_extra_attributes() {
242 assert_eq!(
243 ActiveEnum {
244 enum_name: Alias::new("coinflip_result_type").into_iden(),
245 values: vec!["HEADS", "TAILS"]
246 .into_iter()
247 .map(|variant| Alias::new(variant).into_iden())
248 .collect(),
249 }
250 .impl_active_enum(
251 &WithSerde::None,
252 true,
253 &TokenStream::new(),
254 &bonus_attributes([r#"serde(rename_all = "camelCase")"#]),
255 EntityFormat::Compact,
256 )
257 .to_string(),
258 quote!(
259 #[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum, Copy)]
260 #[sea_orm(
261 rs_type = "String",
262 db_type = "Enum",
263 enum_name = "coinflip_result_type"
264 )]
265 #[serde(rename_all = "camelCase")]
266 pub enum CoinflipResultType {
267 #[sea_orm(string_value = "HEADS")]
268 Heads,
269 #[sea_orm(string_value = "TAILS")]
270 Tails,
271 }
272 )
273 .to_string()
274 );
275 assert_eq!(
276 ActiveEnum {
277 enum_name: Alias::new("coinflip_result_type").into_iden(),
278 values: vec!["HEADS", "TAILS"]
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 &bonus_attributes([r#"serde(rename_all = "camelCase")"#, "ts(export)"]),
288 EntityFormat::Compact,
289 )
290 .to_string(),
291 quote!(
292 #[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum, Copy)]
293 #[sea_orm(
294 rs_type = "String",
295 db_type = "Enum",
296 enum_name = "coinflip_result_type"
297 )]
298 #[serde(rename_all = "camelCase")]
299 #[ts(export)]
300 pub enum CoinflipResultType {
301 #[sea_orm(string_value = "HEADS")]
302 Heads,
303 #[sea_orm(string_value = "TAILS")]
304 Tails,
305 }
306 )
307 .to_string()
308 )
309 }
310
311 #[test]
312 fn test_enum_variant_utf8_encode() {
313 assert_eq!(
314 ActiveEnum {
315 enum_name: Alias::new("ty").into_iden(),
316 values: vec![
317 "Question",
318 "QuestionsAdditional",
319 "Answer",
320 "Other",
321 "/",
322 "//",
323 "A-B-C",
324 "你好",
325 "0/",
326 "0//",
327 "0A-B-C",
328 "0你好",
329 ]
330 .into_iter()
331 .map(|variant| Alias::new(variant).into_iden())
332 .collect(),
333 }
334 .impl_active_enum(
335 &WithSerde::None,
336 true,
337 &TokenStream::new(),
338 &TokenStream::new(),
339 EntityFormat::Compact,
340 )
341 .to_string(),
342 quote!(
343 #[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum, Copy)]
344 #[sea_orm(rs_type = "String", db_type = "Enum", enum_name = "ty")]
345 pub enum Ty {
346 #[sea_orm(string_value = "Question")]
347 Question,
348 #[sea_orm(string_value = "QuestionsAdditional")]
349 QuestionsAdditional,
350 #[sea_orm(string_value = "Answer")]
351 Answer,
352 #[sea_orm(string_value = "Other")]
353 Other,
354 #[sea_orm(string_value = "/")]
355 U002F,
356 #[sea_orm(string_value = "//")]
357 U002FU002F,
358 #[sea_orm(string_value = "A-B-C")]
359 ABC,
360 #[sea_orm(string_value = "你好")]
361 你好,
362 #[sea_orm(string_value = "0/")]
363 _0U002F,
364 #[sea_orm(string_value = "0//")]
365 _0U002FU002F,
366 #[sea_orm(string_value = "0A-B-C")]
367 _0ABC,
368 #[sea_orm(string_value = "0你好")]
369 _0你好,
370 }
371 )
372 .to_string()
373 )
374 }
375}