virtue_next/generate/gen_enum.rs
1use super::{
2 AttributeContainer, Field, FieldBuilder, Impl, ImplFor, Parent, Path, StreamBuilder,
3 StringOrIdent,
4};
5use crate::Result;
6use crate::parse::{Generic, Generics, Visibility};
7use crate::prelude::{Delimiter, Ident, Span, TokenStream};
8
9/// Builder to generate an `enum <Name> { <value> { ... }, ... }`
10///
11/// ```
12/// # use virtue::prelude::Generator;
13/// # let mut generator = Generator::with_name("Fooz");
14/// {
15/// let mut enumgen = generator.generate_enum("Foo");
16/// enumgen
17/// .add_value("ZST")
18/// .make_zst();
19/// enumgen
20/// .add_value("Named")
21/// .add_field("bar", "u16")
22/// .add_field("baz", "String");
23/// enumgen
24/// .add_value("Unnamed")
25/// .make_tuple()
26/// .add_field("", "u16")
27/// .add_field("baz", "String");
28/// }
29/// # generator.assert_eq("enum Foo { ZST , Named { bar : u16 , baz : String , } , Unnamed (u16 , String ,) , }");
30/// # Ok::<_, virtue::Error>(())
31/// ```
32///
33/// Generates:
34/// ```
35/// enum Foo {
36/// ZST,
37/// Named {
38/// bar: u16,
39/// baz: String,
40/// },
41/// Unnamed(u16, String),
42/// };
43/// ```
44pub struct GenEnum<'a, P: Parent> {
45 parent: &'a mut P,
46 name: Ident,
47 visibility: Visibility,
48 generics: Option<Generics>,
49 values: Vec<EnumValue>,
50 derives: Vec<Path>,
51 attributes: Vec<StreamBuilder>,
52 additional: Vec<StreamBuilder>,
53}
54
55impl<'a, P: Parent> GenEnum<'a, P> {
56 pub(crate) fn new(parent: &'a mut P, name: impl Into<String>) -> Self {
57 Self {
58 parent,
59 name: Ident::new(name.into().as_str(), Span::call_site()),
60 visibility: Visibility::Default,
61 generics: None,
62 values: Vec::new(),
63 derives: Vec::new(),
64 attributes: Vec::new(),
65 additional: Vec::new(),
66 }
67 }
68
69 /// Make the enum `pub`. By default the struct will have no visibility modifier and will only be visible in the current scope.
70 pub fn make_pub(&mut self) -> &mut Self {
71 self.visibility = Visibility::Pub;
72 self
73 }
74
75 /// Add a derive macro to the enum.
76 ///
77 /// ```
78 /// # use virtue::prelude::Generator;
79 /// # use virtue::generate::Path;
80 /// # let mut generator = Generator::with_name("Bar");
81 /// generator
82 /// .generate_enum("Foo")
83 /// .with_derive("Clone")
84 /// .with_derive("Default")
85 /// .with_derive(Path::from_iter(vec!["serde", "Deserialize"]));
86 /// # generator.assert_eq("# [derive (Clone , Default , serde ::Deserialize)] enum Foo { }");
87 /// # Ok::<_, virtue::Error>(())
88 /// ```
89 ///
90 /// Generates:
91 /// ```ignore
92 /// #[derive(Clone, Default, serde::Deserialize)]
93 /// enum Foo { }
94 pub fn with_derive(&mut self, derive: impl Into<Path>) -> &mut Self {
95 AttributeContainer::with_derive(self, derive)
96 }
97
98 /// Add derive macros to the enum.
99 ///
100 /// ```
101 /// # use virtue::prelude::Generator;
102 /// # use virtue::generate::Path;
103 /// # let mut generator = Generator::with_name("Bar");
104 /// generator
105 /// .generate_enum("Foo")
106 /// .with_derives([
107 /// "Clone".into(),
108 /// "Default".into(),
109 /// Path::from_iter(vec!["serde", "Deserialize"]),
110 /// ]);
111 /// # generator.assert_eq("# [derive (Clone , Default , serde ::Deserialize)] enum Foo { }");
112 /// # Ok::<_, virtue::Error>(())
113 /// ```
114 ///
115 /// Generates:
116 /// ```ignore
117 /// #[derive(Clone, Default, serde::Deserialize)]
118 /// enum Foo { }
119 pub fn with_derives<T: Into<Path>>(
120 &mut self,
121 derives: impl IntoIterator<Item = T>,
122 ) -> &mut Self {
123 AttributeContainer::with_derives(self, derives)
124 }
125
126 /// Add an attribute to the enum. For `#[derive(...)]`, use [`with_derive`](Self::with_derive)
127 /// instead.
128 ///
129 /// ```
130 /// # use virtue::prelude::Generator;
131 /// # let mut generator = Generator::with_name("Bar");
132 /// generator
133 /// .generate_enum("Foo")
134 /// .with_attribute("serde", |b| {
135 /// b.push_parsed("(untagged)")?;
136 /// Ok(())
137 /// })?;
138 /// # generator.assert_eq("# [serde (untagged)] enum Foo { }");
139 /// # Ok::<_, virtue::Error>(())
140 /// ```
141 ///
142 /// Generates:
143 /// ```ignore
144 /// #[serde(untagged)]
145 /// enum Foo { }
146 /// ```
147 pub fn with_attribute(
148 &mut self,
149 name: impl AsRef<str>,
150 value: impl FnOnce(&mut StreamBuilder) -> Result,
151 ) -> Result<&mut Self> {
152 AttributeContainer::with_attribute(self, name, value)
153 }
154
155 /// Add a parsed attribute to the enum. For `#[derive(...)]`, use [`with_derive`](Self::with_derive)
156 /// instead.
157 ///
158 /// ```
159 /// # use virtue::prelude::Generator;
160 /// # let mut generator = Generator::with_name("Bar");
161 ///
162 /// generator
163 /// .generate_enum("Foo")
164 /// .with_parsed_attribute("serde(untagged)")?;
165 /// # generator.assert_eq("# [serde (untagged)] enum Foo { }");
166 /// # Ok::<_, virtue::Error>(())
167 /// ```
168 ///
169 /// Generates:
170 /// ```ignore
171 /// #[serde(untagged)]
172 /// enum Foo { }
173 /// ```
174 pub fn with_parsed_attribute(&mut self, attribute: impl AsRef<str>) -> Result<&mut Self> {
175 AttributeContainer::with_parsed_attribute(self, attribute)
176 }
177
178 /// Add a token stream as an attribute to the enum. For `#[derive(...)]`, use
179 /// [`with_derive`](Self::with_derive) instead.
180 ///
181 /// ```
182 /// # use virtue::prelude::{Generator, TokenStream};
183 /// # let mut generator = Generator::with_name("Bar");
184 ///
185 /// let attribute = "serde(untagged)".parse::<TokenStream>().unwrap();
186 /// generator
187 /// .generate_enum("Foo")
188 /// .with_attribute_stream(attribute);
189 /// # generator.assert_eq("# [serde (untagged)] enum Foo { }");
190 /// # Ok::<_, virtue::Error>(())
191 /// ```
192 ///
193 /// Generates:
194 /// ```ignore
195 /// #[serde(untagged)]
196 /// enum Foo { }
197 /// ```
198 pub fn with_attribute_stream(&mut self, attribute: impl Into<TokenStream>) -> &mut Self {
199 AttributeContainer::with_attribute_stream(self, attribute)
200 }
201
202 /// Inherit the generic parameters of the parent type.
203 ///
204 /// ```
205 /// # use virtue::prelude::Generator;
206 /// # use virtue::parse::{Generic, Lifetime};
207 /// # use proc_macro2::{Ident, Span};
208 /// # let mut generator = Generator::with_name("Bar").with_lifetime("a");
209 /// // given a derive on enum Bar<'a>
210 /// generator
211 /// .generate_enum("Foo")
212 /// .inherit_generics()
213 /// .add_value("Bar")
214 /// .make_tuple()
215 /// .add_field("bar", "&'a str");
216 /// # generator.assert_eq("enum Foo < 'a > { Bar (&'a str ,) , }");
217 /// # Ok::<_, virtue::Error>(())
218 /// ```
219 ///
220 /// Generates:
221 /// ```ignore
222 /// // given a derive on enum Bar<'a>
223 /// enum Foo<'a> {
224 /// Bar(&'a str)
225 /// }
226 /// ```
227 pub fn inherit_generics(&mut self) -> &mut Self {
228 self.generics = self.parent.generics().cloned();
229 self
230 }
231
232 /// Append generic parameters to the type.
233 ///
234 /// ```
235 /// # use virtue::prelude::Generator;
236 /// # use virtue::parse::{Generic, Lifetime};
237 /// # use proc_macro2::{Ident, Span};
238 /// # let mut generator = Generator::with_name("Bar").with_lifetime("a");
239 /// generator
240 /// .generate_enum("Foo")
241 /// .with_generics([Lifetime { ident: Ident::new("a", Span::call_site()), constraint: vec![] }.into()])
242 /// .add_value("Bar")
243 /// .make_tuple()
244 /// .add_field("bar", "&'a str");
245 /// # generator.assert_eq("enum Foo < 'a > { Bar (&'a str ,) , }");
246 /// # Ok::<_, virtue::Error>(())
247 /// ```
248 ///
249 /// Generates:
250 /// ```ignore
251 /// enum Foo<'a> {
252 /// Bar(&'a str)
253 /// }
254 /// ```
255 pub fn with_generics(&mut self, generics: impl IntoIterator<Item = Generic>) -> &mut Self {
256 self.generics
257 .get_or_insert_with(|| Generics(Vec::new()))
258 .extend(generics);
259 self
260 }
261
262 /// Add a generic parameter to the type.
263 ///
264 /// ```
265 /// # use virtue::prelude::Generator;
266 /// # use virtue::parse::{Generic, Lifetime};
267 /// # use proc_macro2::{Ident, Span};
268 /// # let mut generator = Generator::with_name("Bar").with_lifetime("a");
269 /// generator
270 /// .generate_enum("Foo")
271 /// .with_generic(Lifetime { ident: Ident::new("a", Span::call_site()), constraint: vec![] }.into())
272 /// .add_value("Bar")
273 /// .make_tuple()
274 /// .add_field("bar", "&'a str");
275 /// # generator.assert_eq("enum Foo < 'a > { Bar (&'a str ,) , }");
276 /// # Ok::<_, virtue::Error>(())
277 /// ```
278 ///
279 /// Generates:
280 /// ```ignore
281 /// enum Foo<'a> {
282 /// Bar(&'a str)
283 /// }
284 /// ```
285 pub fn with_generic(&mut self, generic: Generic) -> &mut Self {
286 self.generics
287 .get_or_insert_with(|| Generics(Vec::new()))
288 .push(generic);
289 self
290 }
291
292 /// Add an enum value
293 ///
294 /// Returns a builder for the value that's similar to GenStruct
295 pub fn add_value(&mut self, name: impl Into<String>) -> &mut EnumValue {
296 self.values.push(EnumValue::new(name));
297 self.values.last_mut().unwrap()
298 }
299
300 /// Add an `impl <name> for <enum>`
301 pub fn impl_for(&mut self, name: impl Into<StringOrIdent>) -> ImplFor<'_, Self> {
302 ImplFor::new(self, name.into(), None)
303 }
304
305 /// Generate an `impl <name>` implementation. See [`Impl`] for more information.
306 pub fn r#impl(&mut self) -> Impl<'_, Self> {
307 Impl::with_parent_name(self)
308 }
309
310 /// Generate an `impl <name>` implementation. See [`Impl`] for more information.
311 ///
312 /// Alias for [`impl`] which doesn't need a `r#` prefix.
313 ///
314 /// [`impl`]: #method.impl
315 pub fn generate_impl(&mut self) -> Impl<'_, Self> {
316 Impl::with_parent_name(self)
317 }
318}
319
320impl<P: Parent> AttributeContainer for GenEnum<'_, P> {
321 fn derives(&mut self) -> &mut Vec<Path> {
322 &mut self.derives
323 }
324
325 fn attributes(&mut self) -> &mut Vec<StreamBuilder> {
326 &mut self.attributes
327 }
328}
329
330impl<'a, P: Parent> Parent for GenEnum<'a, P> {
331 fn append(&mut self, builder: StreamBuilder) {
332 self.additional.push(builder);
333 }
334
335 fn name(&self) -> &Ident {
336 &self.name
337 }
338
339 fn generics(&self) -> Option<&Generics> {
340 self.generics.as_ref()
341 }
342
343 fn generic_constraints(&self) -> Option<&crate::parse::GenericConstraints> {
344 None
345 }
346}
347
348impl<'a, P: Parent> Drop for GenEnum<'a, P> {
349 fn drop(&mut self) {
350 let mut builder = StreamBuilder::new();
351 self.build_derives(&mut builder)
352 .build_attributes(&mut builder);
353 if self.visibility == Visibility::Pub {
354 builder.ident_str("pub");
355 }
356 builder
357 .ident_str("enum")
358 .ident(self.name.clone())
359 .append(
360 self.generics()
361 .map(Generics::impl_generics)
362 .unwrap_or_default(),
363 )
364 .group(Delimiter::Brace, |b| {
365 for value in self.values.iter_mut() {
366 build_value(b, value)?;
367 }
368
369 Ok(())
370 })
371 .expect("Could not build enum");
372
373 for additional in std::mem::take(&mut self.additional) {
374 builder.append(additional);
375 }
376 self.parent.append(builder);
377 }
378}
379
380fn build_value(builder: &mut StreamBuilder, value: &mut EnumValue) -> Result {
381 value.build_attributes(builder);
382 builder.ident(value.name.clone());
383
384 match value.value_type {
385 ValueType::Named => builder.group(Delimiter::Brace, |b| {
386 for field in value.fields.iter_mut() {
387 field.build_attributes(b);
388 if field.vis == Visibility::Pub {
389 b.ident_str("pub");
390 }
391 b.ident_str(&field.name)
392 .punct(':')
393 .push_parsed(&field.ty)?
394 .punct(',');
395 }
396 Ok(())
397 })?,
398 ValueType::Unnamed => builder.group(Delimiter::Parenthesis, |b| {
399 for field in value.fields.iter_mut() {
400 field.build_attributes(b);
401 if field.vis == Visibility::Pub {
402 b.ident_str("pub");
403 }
404 b.push_parsed(&field.ty)?.punct(',');
405 }
406 Ok(())
407 })?,
408 ValueType::Zst => builder,
409 };
410
411 builder.punct(',');
412
413 Ok(())
414}
415
416pub struct EnumValue {
417 name: Ident,
418 fields: Vec<Field>,
419 value_type: ValueType,
420 attributes: Vec<StreamBuilder>,
421}
422
423impl EnumValue {
424 fn new(name: impl Into<String>) -> Self {
425 Self {
426 name: Ident::new(name.into().as_str(), Span::call_site()),
427 fields: Vec::new(),
428 value_type: ValueType::Named,
429 attributes: Vec::new(),
430 }
431 }
432
433 /// Make the struct a zero-sized type (no fields)
434 ///
435 /// Any fields will be ignored
436 pub fn make_zst(&mut self) -> &mut Self {
437 self.value_type = ValueType::Zst;
438 self
439 }
440
441 /// Make the struct fields unnamed
442 ///
443 /// The names of any field will be ignored
444 pub fn make_tuple(&mut self) -> &mut Self {
445 self.value_type = ValueType::Unnamed;
446 self
447 }
448
449 /// Add an attribute to the variant.
450 ///
451 /// ```
452 /// # use virtue::prelude::Generator;
453 /// # let mut generator = Generator::with_name("Bar");
454 /// generator
455 /// .generate_enum("Foo")
456 /// .add_value("Bar")
457 /// .with_attribute("serde", |b| {
458 /// b.push_parsed("(rename_all = \"camelCase\")")?;
459 /// Ok(())
460 /// })?;
461 /// # generator.assert_eq("enum Foo { # [serde (rename_all = \"camelCase\")] Bar { } , }");
462 /// # Ok::<_, virtue::Error>(())
463 /// ```
464 ///
465 /// Generates:
466 /// ```ignore
467 /// enum Foo {
468 /// #[serde(rename_all = "camelCase")]
469 /// Bar { }
470 /// }
471 /// ```
472 pub fn with_attribute(
473 &mut self,
474 name: impl AsRef<str>,
475 value: impl FnOnce(&mut StreamBuilder) -> Result,
476 ) -> Result<&mut Self> {
477 AttributeContainer::with_attribute(self, name, value)
478 }
479
480 /// Add a parsed attribute to the variant.
481 ///
482 /// ```
483 /// # use virtue::prelude::Generator;
484 /// # let mut generator = Generator::with_name("Bar");
485 /// generator
486 /// .generate_enum("Foo")
487 /// .add_value("Bar")
488 /// .with_parsed_attribute("serde(rename_all = \"camelCase\")")?;
489 /// # generator.assert_eq("enum Foo { # [serde (rename_all = \"camelCase\")] Bar { } , }");
490 /// # Ok::<_, virtue::Error>(())
491 /// ```
492 ///
493 /// Generates:
494 /// ```ignore
495 /// enum Foo {
496 /// #[serde(rename_all = "camelCase")]
497 /// Bar { }
498 /// }
499 /// ```
500 pub fn with_parsed_attribute(&mut self, attribute: impl AsRef<str>) -> Result<&mut Self> {
501 AttributeContainer::with_parsed_attribute(self, attribute)
502 }
503
504 /// Add a token stream as an attribute to the variant.
505 ///
506 /// ```
507 /// # use virtue::prelude::{Generator, TokenStream};
508 /// # let mut generator = Generator::with_name("Bar");
509 /// let attribute = "serde(rename_all = \"camelCase\")".parse::<TokenStream>().unwrap();
510 /// generator
511 /// .generate_enum("Foo")
512 /// .add_value("Bar")
513 /// .with_attribute_stream(attribute);
514 /// # generator.assert_eq("enum Foo { # [serde (rename_all = \"camelCase\")] Bar { } , }");
515 /// # Ok::<_, virtue::Error>(())
516 /// ```
517 ///
518 /// Generates:
519 /// ```ignore
520 /// enum Foo {
521 /// #[serde(rename_all = "camelCase")]
522 /// Bar { }
523 /// }
524 /// ```
525 pub fn with_attribute_stream(&mut self, attribute: impl Into<TokenStream>) -> &mut Self {
526 AttributeContainer::with_attribute_stream(self, attribute)
527 }
528
529 /// Add a field to the enum value.
530 ///
531 /// Names are ignored when the enum value's fields are unnamed
532 ///
533 /// ```
534 /// # use virtue::prelude::Generator;
535 /// # let mut generator = Generator::with_name("Fooz");
536 /// generator
537 /// .generate_enum("Foo")
538 /// .add_value("Bar")
539 /// .add_field("bar", "u16")
540 /// .add_field("baz", "String");
541 /// # generator.assert_eq("enum Foo { Bar { bar : u16 , baz : String , } , }");
542 /// # Ok::<_, virtue::Error>(())
543 /// ```
544 ///
545 /// Generates:
546 /// ```
547 /// enum Foo {
548 /// Bar {
549 /// bar: u16,
550 /// baz: String
551 /// }
552 /// };
553 /// ```
554 pub fn add_field(
555 &mut self,
556 name: impl Into<String>,
557 ty: impl Into<String>,
558 ) -> FieldBuilder<'_, Self> {
559 let mut fields = FieldBuilder::from(&mut self.fields);
560 fields.add_field(name, ty);
561 fields
562 }
563}
564
565impl AttributeContainer for EnumValue {
566 fn derives(&mut self) -> &mut Vec<Path> {
567 unreachable!("enum variants cannot have derives")
568 }
569
570 fn attributes(&mut self) -> &mut Vec<StreamBuilder> {
571 &mut self.attributes
572 }
573}
574
575enum ValueType {
576 Named,
577 Unnamed,
578 Zst,
579}