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