From df5f46ab2704f2927b1392257b73554eef6d9740 Mon Sep 17 00:00:00 2001 From: "Taro L. Saito" Date: Sat, 11 Nov 2023 22:59:01 -0800 Subject: [PATCH] codec (feature): Support Scala 3 Enum (#3272) Closes #3036 --- .../airframe/codec/Scala3EnumCodecTest.scala | 72 +++++++++++++++++++ .../wvlet/airframe/surface/Surface.scala | 5 ++ .../surface/CompileTimeSurfaceFactory.scala | 39 +++++++++- .../wvlet/airframe/surface/Surfaces.scala | 5 ++ 4 files changed, 120 insertions(+), 1 deletion(-) create mode 100644 airframe-codec/src/test/scala-3/wvlet/airframe/codec/Scala3EnumCodecTest.scala diff --git a/airframe-codec/src/test/scala-3/wvlet/airframe/codec/Scala3EnumCodecTest.scala b/airframe-codec/src/test/scala-3/wvlet/airframe/codec/Scala3EnumCodecTest.scala new file mode 100644 index 0000000000..4c185eea1a --- /dev/null +++ b/airframe-codec/src/test/scala-3/wvlet/airframe/codec/Scala3EnumCodecTest.scala @@ -0,0 +1,72 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package wvlet.airframe.codec + +import wvlet.airframe.json.JSON +import wvlet.airframe.msgpack.spi.MessagePack +import wvlet.airspec.AirSpec + +object Scala3EnumCodecTest extends AirSpec: + enum Color: + case Red, Green, Blue + + test("pack Scala 3 Enum") { + val codec = MessageCodec.of[Color] + Color.values.foreach { c => + val packed = codec.pack(c) + val unpacked = codec.unpack(packed) + unpacked shouldBe c + } + } + + test("handle invalid enum value input") { + val codec = MessageCodec.of[Color] + val msgpack = MessagePack.newBufferPacker.packString("Black").toByteArray + intercept[IllegalArgumentException] { + codec.unpack(msgpack) + } + } + + case class Point(x: Int, y: Int, c: Color) + + test("pack case class with Enum") { + val codec = MessageCodec.of[Point] + val p = Point(1, 2, Color.Red) + val packed = codec.pack(p) + val unpacked = codec.unpack(packed) + unpacked shouldBe p + } + + test("toJson") { + val codec = MessageCodec.of[Point] + val p = Point(1, 2, Color.Red) + val json = codec.toJson(p) + json shouldBe """{"x":1,"y":2,"c":"Red"}""" + } + + case class OptColor(p: Option[Color]) + + test("Option[Enum]") { + val codec = MessageCodec.of[OptColor] + val p = OptColor(Some(Color.Red)) + val json = codec.toJson(p) + json shouldBe """{"p":"Red"}""" + } + + test("Option[Enum] with None") { + val codec = MessageCodec.of[OptColor] + val p = OptColor(None) + val json = codec.toJson(p) + json shouldBe """{}""" + } diff --git a/airframe-surface/src/main/scala-2/wvlet/airframe/surface/Surface.scala b/airframe-surface/src/main/scala-2/wvlet/airframe/surface/Surface.scala index 76640296ce..e6fa9a3aa5 100644 --- a/airframe-surface/src/main/scala-2/wvlet/airframe/surface/Surface.scala +++ b/airframe-surface/src/main/scala-2/wvlet/airframe/surface/Surface.scala @@ -33,6 +33,11 @@ trait Surface extends Serializable { def isMap: Boolean = classOf[Map[_, _]].isAssignableFrom(rawType) def isArray: Boolean = this.isInstanceOf[ArraySurface] + /** + * True if this surface is a Scala3 enum. In Scala 2, it always returns false + */ + def isEnum: Boolean = false + def objectFactory: Option[ObjectFactory] = None def withOuter(outer: AnyRef): Surface = this } diff --git a/airframe-surface/src/main/scala-3/wvlet/airframe/surface/CompileTimeSurfaceFactory.scala b/airframe-surface/src/main/scala-3/wvlet/airframe/surface/CompileTimeSurfaceFactory.scala index ff57cb55c8..e814fa6672 100644 --- a/airframe-surface/src/main/scala-3/wvlet/airframe/surface/CompileTimeSurfaceFactory.scala +++ b/airframe-surface/src/main/scala-3/wvlet/airframe/surface/CompileTimeSurfaceFactory.scala @@ -432,6 +432,44 @@ private[surface] class CompileTimeSurfaceFactory[Q <: Quotes](using quotes: Q) { '{ new GenericSurface(${ clsOf(a1) }, typeArgs = ${ Expr.ofSeq(typeArgs) }.toIndexedSeq) } case r: Refinement => newGenericSurfaceOf(r.info) + case t if t <:< TypeRepr.of[scala.reflect.Enum] && !(t =:= TypeRepr.of[Nothing]) => + /** + * Build a code for finding Enum instance from an input string value: {{ (cl: Class[_], s: String) => + * Try(EnumType.valueOf(s)).toOption }} + */ + val enumType = t.typeSymbol.companionModule + val valueOfMethod = enumType.methodMember("valueOf").headOption match { + case Some(m) => m + case None => + sys.error(s"valueOf method not found in ${t}") + } + val newFn = Lambda( + owner = Symbol.spliceOwner, + tpe = MethodType(List("cl", "s"))( + _ => List(TypeRepr.of[Class[_]], TypeRepr.of[String]), + _ => TypeRepr.of[Option[Any]] + ), + rhsFn = (sym: Symbol, paramRefs: List[Tree]) => { + val strVarRef = paramRefs(1).asExprOf[String].asTerm + val expr: Term = + Select + .unique( + Apply(Select.unique(Ref(t.typeSymbol.companionModule), "valueOf"), List(strVarRef)), + "asInstanceOf" + ).appliedToType(TypeRepr.of[Any]) + val expr2 = ('{ + scala.util.Try(${ expr.asExprOf[Any] }).toOption + }).asExprOf[Option[Any]].asTerm + expr2.changeOwner(sym) + } + ) + + '{ + EnumSurface( + ${ clsOf(t) }, + ${ newFn.asExprOf[(Class[_], String) => Option[Any]] } + ) + } case t if hasStringUnapply(t) => // Build EnumSurface.apply code // EnumSurface(classOf[t], { (cl: Class[_], s: String) => (companion object).unapply(s).asInstanceOf[Option[Any]] } @@ -447,7 +485,6 @@ private[surface] class CompileTimeSurfaceFactory[Q <: Quotes](using quotes: Q) { val strVarRef = paramRefs(1).asExprOf[String].asTerm val expr = Select.unique(Apply(m, List(strVarRef)), "asInstanceOf").appliedToType(TypeRepr.of[Option[Any]]) expr.changeOwner(sym) - } ) '{ diff --git a/airframe-surface/src/main/scala/wvlet/airframe/surface/Surfaces.scala b/airframe-surface/src/main/scala/wvlet/airframe/surface/Surfaces.scala index 7aa00594a9..d950cedf6d 100644 --- a/airframe-surface/src/main/scala/wvlet/airframe/surface/Surfaces.scala +++ b/airframe-surface/src/main/scala/wvlet/airframe/surface/Surfaces.scala @@ -256,6 +256,11 @@ case class OptionSurface(override val rawType: Class[_], elementSurface: Surface } case class JavaEnumSurface(override val rawType: Class[_]) extends GenericSurface(rawType) +/** + * Enum-like surface for Scala 2.x and Scala 3 + * @param rawType + * @param stringExtractor + */ case class EnumSurface(override val rawType: Class[_], stringExtractor: (Class[_], String) => Option[Any]) extends GenericSurface(rawType)