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