From 2f8e87cd888d700a2444c09da0b945b3d686f60b Mon Sep 17 00:00:00 2001
From: TheSchemm <theschemm@gmail.com>
Date: Wed, 4 May 2022 20:36:59 -0500
Subject: [PATCH] Basic support for enum attribute serialization

---
 yaserde/tests/enum.rs                 | 43 ++++++++++++++
 yaserde_derive/src/ser/expand_enum.rs | 82 ++++++++++++++++++++++++++-
 2 files changed, 122 insertions(+), 3 deletions(-)

diff --git a/yaserde/tests/enum.rs b/yaserde/tests/enum.rs
index 19cb5f6..b0822cd 100644
--- a/yaserde/tests/enum.rs
+++ b/yaserde/tests/enum.rs
@@ -147,6 +147,49 @@ fn attribute_enum() {
   deserialize_and_validate!(content, model, XmlStruct);
 }
 
+#[test]
+fn attribute_enum2() {
+  #[derive(YaSerialize)]
+  #[yaserde(rename = "child1")]
+  struct Child1 {
+    #[yaserde(attribute, rename = "val")]
+    pub val: String,
+  }
+
+  #[derive(YaSerialize)]
+  #[yaserde(rename = "child2")]
+  struct Child2 {
+    #[yaserde(attribute)]
+    pub num: u8,
+  }
+
+  #[derive(YaSerialize)]
+  #[yaserde(flatten)]
+  enum Base {
+    #[yaserde(flatten)]
+    C1(Child1),
+  }
+
+  let content = r#"<child1 val="hello world" />"#;
+  let model = Base::C1(Child1 {
+    val: "hello world".into(),
+  });
+  serialize_and_validate!(model, content);
+
+  #[derive(YaSerialize)]
+  #[yaserde(rename = "base")]
+  enum Base2 {
+    #[yaserde(flatten)]
+    C1(Child1),
+  }
+
+  let content = r#"<base><child1 val="hello world" /></base>"#;
+  let model = Base2::C1(Child1 {
+    val: "hello world".into(),
+  });
+  serialize_and_validate!(model, content);
+}
+
 #[test]
 fn unnamed_enum() {
   #[derive(Debug, PartialEq, YaDeserialize, YaSerialize)]
diff --git a/yaserde_derive/src/ser/expand_enum.rs b/yaserde_derive/src/ser/expand_enum.rs
index d739226..ea97fd5 100644
--- a/yaserde_derive/src/ser/expand_enum.rs
+++ b/yaserde_derive/src/ser/expand_enum.rs
@@ -14,11 +14,81 @@ pub fn serialize(
 ) -> TokenStream {
   let inner_enum_inspector = inner_enum_inspector(data_enum, name, root_attributes);
 
+  let get_id = |field: &YaSerdeField| {
+    field
+      .label()
+      .unwrap_or(field.get_type().get_simple_type_visitor())
+  };
+
+  let variant_matches: TokenStream = data_enum
+    .variants
+    .iter()
+    .map(|variant| -> TokenStream {
+      let _attrs = crate::common::YaSerdeAttribute::parse(&variant.attrs);
+
+      let all_fields = variant
+        .fields
+        .iter()
+        .map(|field| YaSerdeField::new(field.clone()));
+
+      let attribute_fields: Vec<_> = all_fields
+        .clone()
+        .into_iter()
+        .filter(|field| {
+          field.is_attribute()
+            || (field.is_flatten() && matches!(field.get_type(), Field::FieldStruct { .. }))
+        })
+        .collect();
+
+      attribute_fields
+        .iter()
+        .map(|field| {
+          let label = variant.ident.clone();
+          let var = get_id(field);
+          let name = name.clone();
+
+          let destructure = if field.get_value_label().is_some() {
+            quote! {{#var, ..}}
+          } else {
+            quote! {(#var, ..)}
+          };
+
+          if field.is_attribute() {
+            quote! { #name::#label { .. } => { }, }
+          } else {
+            match field.get_type() {
+              Field::FieldStruct { .. } => {
+                if root_attributes.flatten {
+                  quote! {
+                    match self {
+                      #name::#label #destructure => {
+                        let (attributes, namespace) = #var.serialize_attributes(
+                          child_attributes,
+                          child_attributes_namespace,
+                        )?;
+                        child_attributes_namespace.extend(&namespace);
+                        child_attributes.extend(attributes);
+                      },
+                      _ => {}
+                    }
+                  }
+                } else {
+                  quote! {}
+                }
+              }
+              _ => quote! { #name::#label { .. } => { },},
+            }
+          }
+        })
+        .collect()
+    })
+    .collect();
+
   implement_serializer(
     name,
     root,
     root_attributes,
-    quote!(),
+    quote!(#variant_matches),
     quote!(match self {
       #inner_enum_inspector
     }),
@@ -38,7 +108,7 @@ fn inner_enum_inspector(
 
       let label = &variant.ident;
       let label_name = build_label_name(label, &variant_attrs, &root_attributes.default_namespace);
-
+      
       match variant.fields {
         Fields::Unit => quote! {
           &#name::#label => {
@@ -205,7 +275,13 @@ fn inner_enum_inspector(
                     }
                   })
                 }
-                Field::FieldStruct { .. } => write_element(&match_field(&serialize)),
+                Field::FieldStruct { .. } => {
+                  if variant_attrs.flatten || field.is_flatten() {
+                     match_field(&quote!{ ::yaserde::YaSerialize::serialize(item, writer)?})
+                   } else {
+                     write_element(&match_field(&serialize))
+                   }
+                }
                 Field::FieldString => match_field(&write_element(&write_string_chars)),
                 _simple_type => match_field(&write_simple_type),
               }