feat: Handle contextual map keys #114

avro-kotlin · Jan 28, 2024 · dd0eb34 · dd0eb34
1 parent 563c20f
commit dd0eb34
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 13 deletions.
diff --git a/src/main/kotlin/com/github/avrokotlin/avro4k/schema/SchemaFor.kt b/src/main/kotlin/com/github/avrokotlin/avro4k/schema/SchemaFor.kt
@@ -112,19 +112,21 @@ class MapSchemaFor(
 ) : SchemaFor {
     override fun schema(): Schema {
         val keyType = descriptor.getElementDescriptor(0).unwrapValueClass
-        if (keyType.kind !is PrimitiveKind && keyType.kind != SerialKind.ENUM) {
-            throw RuntimeException("Avro4k only supports primitive and enum kinds as the map key. Actual: ${descriptor.getElementDescriptor(0)}")
+            .let { if (it.kind == SerialKind.CONTEXTUAL) serializersModule.getContextualDescriptor(it)?.unwrapValueClass else it }
+        if (keyType != null) {
+            if (keyType.kind is PrimitiveKind || keyType.kind == SerialKind.ENUM) {
+                val valueSchema =
+                    schemaFor(
+                        serializersModule,
+                        descriptor.getElementDescriptor(1),
+                        descriptor.getElementAnnotations(1),
+                        configuration,
+                        resolvedSchemas
+                    ).schema()
+                return Schema.createMap(valueSchema)
+            }
         }
-
-        val valueSchema =
-            schemaFor(
-                serializersModule,
-                descriptor.getElementDescriptor(1),
-                descriptor.getElementAnnotations(1),
-                configuration,
-                resolvedSchemas
-            ).schema()
-        return Schema.createMap(valueSchema)
+        throw RuntimeException("Avro4k only supports primitive and enum kinds as the map key. Actual: ${descriptor.getElementDescriptor(0)}")
     }
 }
 

diff --git a/src/test/kotlin/com/github/avrokotlin/avro4k/endecode/MapEncoderTest.kt b/src/test/kotlin/com/github/avrokotlin/avro4k/endecode/MapEncoderTest.kt
@@ -1,11 +1,21 @@
 package com.github.avrokotlin.avro4k.endecode
 
+import com.github.avrokotlin.avro4k.Avro
 import com.github.avrokotlin.avro4k.record
 import io.kotest.core.factory.TestFactory
 import io.kotest.core.spec.style.StringSpec
 import io.kotest.core.spec.style.stringSpec
+import kotlinx.serialization.Contextual
+import kotlinx.serialization.ExperimentalSerializationApi
+import kotlinx.serialization.InternalSerializationApi
+import kotlinx.serialization.KSerializer
 import kotlinx.serialization.SerialName
 import kotlinx.serialization.Serializable
+import kotlinx.serialization.descriptors.PrimitiveKind
+import kotlinx.serialization.descriptors.buildSerialDescriptor
+import kotlinx.serialization.encoding.Decoder
+import kotlinx.serialization.encoding.Encoder
+import kotlinx.serialization.modules.serializersModuleOf
 import java.nio.ByteBuffer
 
 class MapEncoderTest : StringSpec({
@@ -41,6 +51,16 @@ fun mapEncoderTests(enDecoder: EnDecoder): TestFactory {
             )
         }
 
+        "encode/decode a Map<non serializable key, String>" {
+            @Serializable
+            data class StringStringTest(val a: Map<@Contextual NonSerializableKey, String>)
+            enDecoder.avro = Avro(serializersModule = serializersModuleOf(NonSerializableKey::class, NonSerializableKeyKSerializer()))
+            enDecoder.testEncodeDecode(
+                StringStringTest(mapOf(NonSerializableKey("a") to "x", NonSerializableKey("b") to "y", NonSerializableKey("c") to "z")),
+                record(mapOf("a" to "x", "b" to "y", "c" to "z"))
+            )
+        }
+
         "encode/decode a Map<int value class, String>" {
             @Serializable
             data class StringStringTest(val a: Map<MapIntKey, String>)
@@ -114,4 +134,18 @@ private enum class MyEnum {
 
     @SerialName("z")
     C,
-}
+}
+
+data class NonSerializableKey(val value: String)
+
+@OptIn(ExperimentalSerializationApi::class)
+class NonSerializableKeyKSerializer : KSerializer<NonSerializableKey> {
+    @OptIn(InternalSerializationApi::class)
+    override val descriptor = buildSerialDescriptor("NonSerializableKey", PrimitiveKind.STRING)
+
+    override fun deserialize(decoder: Decoder) = NonSerializableKey(decoder.decodeString())
+
+    override fun serialize(encoder: Encoder, value: NonSerializableKey) {
+        encoder.encodeString(value.value)
+    }
+}