Add DataTemplates support for Unions

coursera · Jan 16, 2018 · 485fb10 · 485fb10
1 parent 9bb9bb0
commit 485fb10
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 18 deletions.
diff --git a/scala/generator-test/src/test/scala/org/coursera/courier/generator/UnionGeneratorTest.scala b/scala/generator-test/src/test/scala/org/coursera/courier/generator/UnionGeneratorTest.scala
@@ -165,6 +165,6 @@ class UnionGeneratorTest extends GeneratorTest with SchemaFixtures {
     val original = WithPrimitiveCustomTypesUnion(origMember)
     val json = DataTemplates.writeUnion(origMember)
     val roundTripped = DataTemplates.readUnion[WithPrimitiveCustomTypesUnion.Union](json)
-    assert(original === roundTripped)
+    assert(origMember === roundTripped)
   }
 }
diff --git a/scala/runtime/src/main/scala/org/coursera/courier/templates/DataTemplates.scala b/scala/runtime/src/main/scala/org/coursera/courier/templates/DataTemplates.scala
@@ -213,28 +213,36 @@ object DataTemplates {
   private[this] def newDataMapTemplate[T <: DataTemplate[DataMap]](
       data: DataMap, clazz: Class[T]): T = {
     val companionInstance = companion(clazz)
-    val companionClass = companionInstance.getClass
+    val applyMethod = getDataMapConstructor(companionInstance)
+
+    applyMethod.invoke(companionInstance, data, DataConversion.SetReadOnly).asInstanceOf[T]
+  }
 
-    // Reflect for the correct method to instantiate the record from a DataMap and DataConversion.
-    // The correct method is usually "build", since pre-2.0.3. Before then it was "apply".
-    // In case this method is reading templates generated in an old version, we start with
-    // "build" and fall back on "apply".
-    def getDataMapConstructor(methodName: String): Method = {
+  /**
+    * Reflect for the correct method to instantiate the record from a DataMap and DataConversion.
+    * The correct method is usually "build", since pre-2.0.3. Before then it was "apply".
+    * In case this method is reading templates generated in an old version, we start with
+    * "build" and fall back on "apply".
+    *
+    * @param tmplCompanionInstance the result of calling the private method companion(clazz)
+    * @return a method that can be construct a Union or Record template when
+    *         invoked a DataMap and DataConversion as args
+    */
+  private[this] def getDataMapConstructor(tmplCompanionInstance: AnyRef): Method = {
+    val companionClass = tmplCompanionInstance.getClass
+    def getWithName(methodName: String): Method = {
       companionClass.getDeclaredMethod(
         methodName,
         classOf[DataMap],
         classOf[DataConversion]
       )
     }
 
-    val applyMethod = try {
-      getDataMapConstructor(methodName = "build")
+    try {
+      getWithName("build")
     } catch {
-      case e: NoSuchMethodException =>
-        getDataMapConstructor(methodName = "apply")
+      case e: NoSuchMethodException => getWithName("apply")
     }
-
-    applyMethod.invoke(companionInstance, data, DataConversion.SetReadOnly).asInstanceOf[T]
   }
 
   private[this] def newDataListTemplate[T <: DataTemplate[DataList]](
@@ -250,11 +258,8 @@ object DataTemplates {
 
   private[this] def newUnionTemplate[T <: UnionTemplate](data: DataMap, clazz: Class[T]): T = {
     val companionInstance = companion(clazz)
-    val applyMethod =
-      companionInstance.getClass.getDeclaredMethod(
-        "apply",
-        classOf[DataMap],
-        classOf[DataConversion])
+    val applyMethod = getDataMapConstructor(companionInstance)
+
     applyMethod.invoke(companionInstance, data, DataConversion.SetReadOnly).asInstanceOf[T]
   }
 }