add support for transformed primitve map keys (#627)

* add support for transformed primitve map keys * run sbt prepare * restore formatting annotation * rename for clarity
zio · Dec 27, 2023 · 830017e · 830017e
1 parent 3cf1987
commit 830017e
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 2 deletions.
diff --git a/zio-schema-json/shared/src/main/scala/zio/schema/codec/JsonCodec.scala b/zio-schema-json/shared/src/main/scala/zio/schema/codec/JsonCodec.scala
@@ -285,11 +285,26 @@ object JsonCodec {
       }
     //scalafmt: { maxColumn = 120, optIn.configStyleArguments = true }
 
+    private[codec] def transformFieldEncoder[A, B](
+      schema: Schema[A],
+      g: B => Either[String, A]
+    ): Option[JsonFieldEncoder[B]] =
+      jsonFieldEncoder(schema).map { fieldEncoder =>
+        new JsonFieldEncoder[B] {
+          override def unsafeEncodeField(b: B): String =
+            g(b) match {
+              case Left(_)  => throw new RuntimeException(s"Failed to encode field $b")
+              case Right(a) => fieldEncoder.unsafeEncodeField(a)
+            }
+        }
+      }
+
     private[codec] def jsonFieldEncoder[A](schema: Schema[A]): Option[JsonFieldEncoder[A]] =
       schema match {
         case Schema.Primitive(StandardType.StringType, _) => Option(JsonFieldEncoder.string)
         case Schema.Primitive(StandardType.LongType, _)   => Option(JsonFieldEncoder.long)
         case Schema.Primitive(StandardType.IntType, _)    => Option(JsonFieldEncoder.int)
+        case Schema.Transform(c, _, g, a, _)              => transformFieldEncoder(a.foldLeft(c)((s, a) => s.annotate(a)), g)
         case Schema.Lazy(inner)                           => jsonFieldEncoder(inner())
         case _                                            => None
       }
@@ -617,8 +632,10 @@ object JsonCodec {
         case Schema.Primitive(StandardType.StringType, _) => Option(JsonFieldDecoder.string)
         case Schema.Primitive(StandardType.LongType, _)   => Option(JsonFieldDecoder.long)
         case Schema.Primitive(StandardType.IntType, _)    => Option(JsonFieldDecoder.int)
-        case Schema.Lazy(inner)                           => jsonFieldDecoder(inner())
-        case _                                            => None
+        case Schema.Transform(c, f, _, a, _) =>
+          jsonFieldDecoder(a.foldLeft(c)((s, a) => s.annotate(a))).map(_.mapOrFail(f))
+        case Schema.Lazy(inner) => jsonFieldDecoder(inner())
+        case _                  => None
       }
 
     private def dynamicDecoder(schema: Schema.Dynamic): ZJsonDecoder[DynamicValue] = {

diff --git a/zio-schema-json/shared/src/test/scala-2/zio/schema/codec/JsonCodecSpec.scala b/zio-schema-json/shared/src/test/scala-2/zio/schema/codec/JsonCodecSpec.scala
@@ -163,6 +163,19 @@ object JsonCodecSpec extends ZIOSpecDefault {
             """{"0":{"first":0,"second":true},"1":{"first":1,"second":false}}"""
           )
         )
+      },
+      test("of complex keys with transformation to primitive keys") {
+        assertEncodes(
+          Schema
+            .map[KeyWrapper, ValueWrapper],
+          Map(
+            KeyWrapper("wrapped_key_1") -> ValueWrapper(value = "some_value"),
+            KeyWrapper("wrapped_key_2") -> ValueWrapper(value = "some_other_value")
+          ),
+          charSequenceToByteChunk(
+            """{"wrapped_key_1":{"value":"some_value"},"wrapped_key_2":{"value":"some_other_value"}}"""
+          )
+        )
       }
     ),
     suite("Set")(
@@ -617,6 +630,19 @@ object JsonCodecSpec extends ZIOSpecDefault {
             """{"0":{"first":0,"second":true},"1":{"first":1,"second":false}}"""
           )
         )
+      },
+      test("of primitive keys with transformation to complex keys") {
+        assertDecodes(
+          Schema
+            .map[KeyWrapper, ValueWrapper],
+          Map(
+            KeyWrapper("wrapped_key_1") -> ValueWrapper(value = "some_value"),
+            KeyWrapper("wrapped_key_2") -> ValueWrapper(value = "some_other_value")
+          ),
+          charSequenceToByteChunk(
+            """{"wrapped_key_1":{"value":"some_value"},"wrapped_key_2":{"value":"some_other_value"}}"""
+          )
+        )
       }
     ),
     suite("zio.json.ast.Json decoding")(
@@ -1605,4 +1631,23 @@ object JsonCodecSpec extends ZIOSpecDefault {
   object ListAndMap {
     implicit lazy val schema: Schema[ListAndMap] = DeriveSchema.gen[ListAndMap]
   }
+
+  final case class KeyWrapper(key: String)
+
+  object KeyWrapper {
+    implicit lazy val schema: Schema[KeyWrapper] = Schema[String].transform(KeyWrapper.apply, _.key)
+  }
+
+  final case class ValueWrapper(value: String)
+
+  object ValueWrapper {
+    implicit lazy val schema: Schema[ValueWrapper] = DeriveSchema.gen[ValueWrapper]
+  }
+
+  final case class MapOfComplexKeysAndValues(map: Map[KeyWrapper, ValueWrapper])
+
+  object MapOfComplexKeysAndValues {
+    implicit lazy val mapSchema: Schema[Map[KeyWrapper, ValueWrapper]] = Schema.map[KeyWrapper, ValueWrapper]
+    implicit lazy val schema: Schema[MapOfComplexKeysAndValues]        = DeriveSchema.gen[MapOfComplexKeysAndValues]
+  }
 }