Add method to add attributes as TokenStreams and missing with_parsed_…

…attribute
bincode-org · Jun 13, 2024 · 55ef8d1 · 55ef8d1
1 parent b0e17f1
commit 55ef8d1
Show file tree

Hide file tree

Showing 3 changed files with 135 additions and 3 deletions.
diff --git a/src/generate/gen_enum.rs b/src/generate/gen_enum.rs
@@ -3,7 +3,7 @@ use super::{
     StringOrIdent,
 };
 use crate::parse::{Generic, Generics, Visibility};
-use crate::prelude::{Delimiter, Ident, Span};
+use crate::prelude::{Delimiter, Ident, Span, TokenStream};
 use crate::Result;
 
 /// Builder to generate an `enum <Name> { <value> { ... }, ... }`
@@ -152,6 +152,53 @@ impl<'a, P: Parent> GenEnum<'a, P> {
         AttributeContainer::with_attribute(self, name, value)
     }
 
+    /// Add a parsed attribute to the enum. For `#[derive(...)]`, use [`with_derive`](Self::with_derive)
+    /// instead.
+    ///
+    /// ```
+    /// # use virtue::prelude::Generator;
+    /// # let mut generator = Generator::with_name("Bar");
+    ///
+    /// generator
+    ///     .generate_enum("Foo")
+    ///     .with_parsed_attribute("serde(untagged)")?;
+    /// # generator.assert_eq("# [serde (untagged)] enum Foo { }");
+    /// # Ok::<_, virtue::Error>(())
+    /// ```
+    ///
+    /// Generates:
+    /// ```ignore
+    /// #[serde(untagged)]
+    /// enum Foo { }
+    /// ```
+    pub fn with_parsed_attribute(&mut self, attribute: impl AsRef<str>) -> Result<&mut Self> {
+        AttributeContainer::with_parsed_attribute(self, attribute)
+    }
+
+    /// Add a token stream as an attribute to the enum. For `#[derive(...)]`, use
+    /// [`with_derive`](Self::with_derive) instead.
+    ///
+    /// ```
+    /// # use virtue::prelude::{Generator, TokenStream};
+    /// # let mut generator = Generator::with_name("Bar");
+    ///
+    /// let attribute = "serde(untagged)".parse::<TokenStream>().unwrap();
+    /// generator
+    ///     .generate_enum("Foo")
+    ///     .with_attribute_stream(attribute);
+    /// # generator.assert_eq("# [serde (untagged)] enum Foo { }");
+    /// # Ok::<_, virtue::Error>(())
+    /// ```
+    ///
+    /// Generates:
+    /// ```ignore
+    /// #[serde(untagged)]
+    /// enum Foo { }
+    /// ```
+    pub fn with_attribute_stream(&mut self, attribute: impl Into<TokenStream>) -> &mut Self {
+        AttributeContainer::with_attribute_stream(self, attribute)
+    }
+
     /// Inherit the generic parameters of the parent type.
     ///
     /// ```
@@ -454,6 +501,31 @@ impl EnumValue {
         AttributeContainer::with_parsed_attribute(self, attribute)
     }
 
+    /// Add a token stream as an attribute to the variant.
+    ///
+    /// ```
+    /// # use virtue::prelude::{Generator, TokenStream};
+    /// # let mut generator = Generator::with_name("Bar");
+    /// let attribute = "serde(rename_all = \"camelCase\")".parse::<TokenStream>().unwrap();
+    /// generator
+    ///     .generate_enum("Foo")
+    ///     .add_value("Bar")
+    ///     .with_attribute_stream(attribute);
+    /// # generator.assert_eq("enum Foo { # [serde (rename_all = \"camelCase\")] Bar { } , }");
+    /// # Ok::<_, virtue::Error>(())
+    /// ```
+    ///
+    /// Generates:
+    /// ```ignore
+    /// enum Foo {
+    ///     #[serde(rename_all = "camelCase")]
+    ///     Bar { }
+    /// }
+    /// ```
+    pub fn with_attribute_stream(&mut self, attribute: impl Into<TokenStream>) -> &mut Self {
+        AttributeContainer::with_attribute_stream(self, attribute)
+    }
+
     /// Add a field to the enum value.
     ///
     /// Names are ignored when the enum value's fields are unnamed

diff --git a/src/generate/gen_struct.rs b/src/generate/gen_struct.rs
@@ -3,7 +3,7 @@ use super::{
     StringOrIdent,
 };
 use crate::parse::{Generic, Generics, Visibility};
-use crate::prelude::{Delimiter, Ident, Span};
+use crate::prelude::{Delimiter, Ident, Span, TokenStream};
 use crate::Result;
 
 /// Builder to generate a struct.
@@ -277,6 +277,31 @@ impl<'a, P: Parent> GenStruct<'a, P> {
         AttributeContainer::with_parsed_attribute(self, attribute)
     }
 
+    /// Add a token stream as an attribute to the struct. For `#[derive(...)]`, use
+    /// [`with_derive`](Self::with_derive) instead.
+    ///
+    /// ```
+    /// # use virtue::prelude::{Generator, TokenStream};
+    /// # use std::str::FromStr;
+    /// # let mut generator = Generator::with_name("Bar");
+    ///
+    /// let attribute = "serde(rename_all = \"camelCase\")".parse::<TokenStream>().unwrap();
+    /// generator
+    ///     .generate_struct("Foo")
+    ///     .with_attribute_stream(attribute);
+    /// # generator.assert_eq("# [serde (rename_all = \"camelCase\")] struct Foo { }");
+    /// # Ok::<_, virtue::Error>(())
+    /// ```
+    ///
+    /// Generates:
+    /// ```ignore
+    /// #[serde(rename_all = "camelCase")]
+    /// struct Foo { }
+    /// ```
+    pub fn with_attribute_stream(&mut self, attribute: impl Into<TokenStream>) -> &mut Self {
+        AttributeContainer::with_attribute_stream(self, attribute)
+    }
+
     /// Add a field to the struct.
     ///
     /// Names are ignored when the Struct's fields are unnamed

diff --git a/src/generate/mod.rs b/src/generate/mod.rs
@@ -24,7 +24,7 @@ mod stream_builder;
 use crate::parse::Visibility;
 use crate::{
     parse::{GenericConstraints, Generics},
-    prelude::{Delimiter, Ident},
+    prelude::{Delimiter, Ident, TokenStream},
 };
 use std::fmt;
 use std::marker::PhantomData;
@@ -258,6 +258,33 @@ impl<P> FieldBuilder<'_, P> {
         Ok(self)
     }
 
+    /// Add a token stream as an attribute to the field.
+    ///
+    /// ```
+    /// # use virtue::prelude::{Generator, TokenStream};
+    /// # let mut generator = Generator::with_name("Fooz");
+    /// let attribute = "serde(default)".parse::<TokenStream>().unwrap();
+    /// generator
+    ///     .generate_struct("Foo")
+    ///     .add_field("foo", "u16")
+    ///     .make_pub()
+    ///     .with_attribute_stream(attribute);
+    /// # generator.assert_eq("struct Foo { # [serde (default)] pub foo : u16 , }");
+    /// # Ok::<_, virtue::Error>(())
+    /// ```
+    ///
+    /// Generates:
+    /// ```ignore
+    /// struct Foo {
+    ///     #[serde(default)]
+    ///     pub bar: u16
+    /// }
+    /// ```
+    pub fn with_attribute_stream(&mut self, attribute: impl Into<TokenStream>) -> &mut Self {
+        self.current().with_attribute_stream(attribute);
+        self
+    }
+
     /// Add a field to the parent type.
     ///
     /// ```
@@ -342,6 +369,14 @@ trait AttributeContainer {
         Ok(self)
     }
 
+    fn with_attribute_stream(&mut self, attribute: impl Into<TokenStream>) -> &mut Self {
+        let stream = StreamBuilder {
+            stream: attribute.into(),
+        };
+        self.attributes().push(stream);
+        self
+    }
+
     fn build_derives(&mut self, b: &mut StreamBuilder) -> &mut Self {
         let derives = std::mem::take(self.derives());
         if !derives.is_empty() {