From 077d69ea8b69f00f6b0ea17053bfb199e9353a1b Mon Sep 17 00:00:00 2001
From: Kyle Davis <github@kyle.dev>
Date: Mon, 18 Mar 2024 20:09:39 -0700
Subject: [PATCH] fix: update enum serialization (#2)

---
 src/model.rs  | 32 +++++++++++++++++++++++++-----
 tests/main.rs | 54 +++++++++++++++++++++++++++++++++++++++++++++++----
 2 files changed, 77 insertions(+), 9 deletions(-)

diff --git a/src/model.rs b/src/model.rs
index 3dbf04c..406fe75 100644
--- a/src/model.rs
+++ b/src/model.rs
@@ -109,9 +109,9 @@ impl ConvexField {
     let ignore_attributes = quote! {
       #[allow(non_snake_case)]
     };
-    let struct_attributes = quote! {
-
-      #[derive(serde::Deserialize, serde::Serialize, Clone, Debug)]
+    // Note: We need a custom serialize to avoid unions printing as objects.
+    let enum_struct_attributes = quote! {
+      #[derive(::serde::Deserialize, Clone, Debug)]
     };
 
     match &self.t {
@@ -126,6 +126,7 @@ impl ConvexField {
         let mut enum_kinds = Vec::new();
         let mut extract_arms = Vec::new();
         let mut json_arms: Vec<TokenStream> = Vec::new();
+        let mut serialize_arms: Vec<TokenStream> = Vec::new();
         let mut i = 0;
         for t in types {
           i += 1;
@@ -146,6 +147,9 @@ impl ConvexField {
                 json_arms.push(quote! {
                   | #struct_name::#branch_name => ::serde_json::Value::Null,
                 });
+                serialize_arms.push(quote! {
+                  | #struct_name::#branch_name => ().serialize(serializer),
+                });
               } else {
                 enum_kinds.push(quote! {
                   #branch_name(#branch_type),
@@ -153,6 +157,9 @@ impl ConvexField {
                 json_arms.push(quote! {
                   | #struct_name::#branch_name(value) => ::serde_json::json!(value),
                 });
+                serialize_arms.push(quote! {
+                  | #struct_name::#branch_name(ref value) => value.serialize(serializer),
+                });
               }
             },
             | None => {
@@ -162,6 +169,9 @@ impl ConvexField {
               json_arms.push(quote! {
                 | #struct_name::#branch_name(value) => ::serde_json::json!(value),
               });
+              serialize_arms.push(quote! {
+                | #struct_name::#branch_name(ref value) => value.serialize(serializer),
+              });
             },
           };
           match t {
@@ -241,12 +251,24 @@ impl ConvexField {
 
         structs.push(quote! {
           #ignore_attributes
-          #struct_attributes
+          #enum_struct_attributes
           pub enum #struct_name {
             #( #enum_kinds )*
           }
         });
 
+        impls.push(quote! {
+          #ignore_attributes
+          impl ::serde::Serialize for #struct_name {
+            fn serialize<S>(&self, serializer: S) -> ::core::result::Result<S::Ok, S::Error>
+            where S: ::serde::Serializer {
+              match *self {
+                #( #serialize_arms )*
+              }
+            }
+          }
+        });
+
         impls.push(quote! {
           #ignore_attributes
           impl #struct_name {
@@ -312,7 +334,7 @@ impl ConvexField {
       #[allow(non_snake_case)]
     };
     let struct_attributes = quote! {
-      #[derive(serde::Deserialize, serde::Serialize, Clone, Debug)]
+      #[derive(::serde::Serialize, ::serde::Deserialize, Clone, Debug)]
     };
     let mut structs = Vec::new();
     let mut rendered_fields = Vec::new();
diff --git a/tests/main.rs b/tests/main.rs
index 95419bf..2be18fb 100644
--- a/tests/main.rs
+++ b/tests/main.rs
@@ -59,10 +59,56 @@ fn example() {
     "seven": "seven",
     "eight": false,
     "nine": 9,
-    // TODO: Fix this.
-    "ten": {
-      "Variant2": 10.0
-    },
+    "ten": 10.0,
+  });
+
+  let actual_json_data = json!(model);
+  assert_eq!(expected_json_data, actual_json_data);
+}
+
+convex_model!(BasicUnion { value: v.union(v.string(), v.number()) });
+
+#[test]
+fn example2() {
+  let convex_data = Value::Object(btreemap! {
+    "value".into() => Value::String("hi".into()),
+  });
+
+  let model = BasicUnion::from_convex_value(&convex_data)
+    .expect("Model should parse data");
+
+  if let BasicUnionValue::Variant1(value) = &model.value {
+    assert_eq!("hi", value);
+  } else {
+    panic!("Expected 10.0")
+  }
+
+  let expected_json_data = json!({
+    "value": "hi",
+  });
+
+  let actual_json_data = json!(model);
+  assert_eq!(expected_json_data, actual_json_data);
+}
+
+convex_model!(NullUnion { value: v.union(v.string(), v.null()) });
+
+#[test]
+fn example3() {
+  let convex_data = Value::Object(btreemap! {
+    "value".into() => Value::Null,
+  });
+
+  let model = NullUnion::from_convex_value(&convex_data)
+    .expect("Model should parse data");
+
+  if let NullUnionValue::Variant2 = &model.value {
+  } else {
+    panic!("Expected Unit")
+  }
+
+  let expected_json_data = json!({
+    "value": null,
   });
 
   let actual_json_data = json!(model);