Merge pull request #238 from alexjg/fix-options-of-sealed-traits-enco…

…ding Fix an error where options of sealed traits weren't encoded correctly
sksamuel · Oct 11, 2018 · c6053d8 · c6053d8
2 parents 9bea085 + 44000a0
commit c6053d8
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 3 deletions.
diff --git a/avro4s-core/src/test/scala/com/sksamuel/avro4s/record/encoder/SealedTraitEncoderTest.scala b/avro4s-core/src/test/scala/com/sksamuel/avro4s/record/encoder/SealedTraitEncoderTest.scala
@@ -54,6 +54,13 @@ class SealedTraitEncoderTest extends FunSuite with Matchers {
     Encoder[Dibble].encode(Dobble, schema) shouldBe new GenericData.EnumSymbol(schema, Dobble)
     Encoder[Dibble].encode(Dabble, schema) shouldBe new GenericData.EnumSymbol(schema, Dabble)
   }
+
+  test("options of sealed traits should be encoded correctly") {
+    val schema = AvroSchema[MeasurableThing]
+    val record = Encoder[MeasurableThing].encode(MeasurableThing(Some(WidthDimension(1.23))), schema).asInstanceOf[GenericRecord]
+    val width = record.get("dimension").asInstanceOf[GenericRecord]
+    width.get("width") shouldBe 1.23
+  }
 }
 
 sealed trait Dibble
@@ -75,3 +82,9 @@ sealed trait Nibble
 case class Nobble(str: String, place: String) extends Nibble
 case class Nabble(str: String, age: Int) extends Nibble
 case class Napper(nibble: Nibble)
+
+sealed trait Dimension
+case class HeightDimension(height: Double) extends Dimension
+case class WidthDimension(width: Double) extends Dimension
+case class MeasurableThing(dimension: Option[Dimension])
+
diff --git a/avro4s-macros/src/main/scala/com/sksamuel/avro4s/Encoder.scala b/avro4s-macros/src/main/scala/com/sksamuel/avro4s/Encoder.scala
@@ -205,9 +205,15 @@ object Encoder extends CoproductEncoders with TupleEncoders {
     import scala.collection.JavaConverters._
 
     override def encode(t: Option[T], schema: Schema): AnyRef = {
-      // if the option is none we just return null, otherwise we encode the value
-      // by finding the non null schema
-      val nonNullSchema = schema.getTypes.asScala.find(_.getType != Schema.Type.NULL).get
+      val nonNullSchema = schema.getTypes().size match {
+        // if the option is none we just return null, otherwise we encode the value
+        // by finding the non null schema
+        case 2 => schema.getTypes.asScala.find(_.getType != Schema.Type.NULL).get
+        case otherwise => {
+          val remainingSchemas = schema.getTypes.asScala.filter(_.getType != Schema.Type.NULL)
+          Schema.createUnion(remainingSchemas.toList.asJava)
+        }
+      }
       t.map(encoder.encode(_, nonNullSchema)).orNull
     }
   }