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-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] = {

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]
+ }
 }