diff --git a/src/struct.rs b/src/struct.rs index 8feae8a..05f7010 100644 --- a/src/struct.rs +++ b/src/struct.rs @@ -14,6 +14,9 @@ pub struct Struct { /// Struct fields fields: Fields, + + /// The attributes for this struct. + attributes: Vec, } impl Struct { @@ -22,6 +25,7 @@ impl Struct { Struct { type_def: TypeDef::new(name), fields: Fields::Empty, + attributes: vec![], } } @@ -63,6 +67,12 @@ impl Struct { self } + /// Adds an attribute to the struct (e.g. `"#[some_attribute]"`) + pub fn attribute(&mut self, attribute: &str) -> &mut Self { + self.attributes.push(attribute.to_string()); + self + } + /// Specify lint attribute to supress a warning or error. pub fn allow(&mut self, allow: &str) -> &mut Self { self.type_def.allow(allow); @@ -110,6 +120,9 @@ impl Struct { /// Formats the struct using the given formatter. pub fn fmt(&self, fmt: &mut Formatter<'_>) -> fmt::Result { + for m in self.attributes.iter() { + write!(fmt, "{}\n", m)?; + } self.type_def.fmt_head("struct", &[], fmt)?; self.fields.fmt(fmt)?; diff --git a/tests/codegen.rs b/tests/codegen.rs index 77e7346..a44e3b5 100644 --- a/tests/codegen.rs +++ b/tests/codegen.rs @@ -41,6 +41,44 @@ struct Foo { assert_eq!(scope.to_string(), &expect[1..]); } +#[test] +fn struct_with_attribute() { + let mut scope = Scope::new(); + let mut struct_ = Struct::new("Foo"); + let field = Field::new("one", "usize"); + struct_.push_field(field); + struct_.attribute("#[test]"); + scope.push_struct(struct_); + + let expect = r#" +#[test] +struct Foo { + one: usize, +}"#; + + assert_eq!(scope.to_string(), &expect[1..]); +} + +#[test] +fn struct_with_multiple_attributes() { + let mut scope = Scope::new(); + let mut struct_ = Struct::new("Foo"); + let field = Field::new("one", "usize"); + struct_.push_field(field); + struct_.attribute("#[test]"); + struct_.attribute("#[cfg(target_os = \"linux\")]"); + scope.push_struct(struct_); + + let expect = r#" +#[test] +#[cfg(target_os = "linux")] +struct Foo { + one: usize, +}"#; + + assert_eq!(scope.to_string(), &expect[1..]); +} + #[test] fn single_struct_documented_field() { let mut scope = Scope::new();