Skip to content

Commit

Permalink
codec (feature): Support Scala 3 Enum (#3272)
Browse files Browse the repository at this point in the history
Closes #3036
  • Loading branch information
xerial committed Nov 12, 2023
1 parent 188e2e9 commit df5f46a
Show file tree
Hide file tree
Showing 4 changed files with 120 additions and 1 deletion.
Original file line number Diff line number Diff line change
@@ -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 """{}"""
}
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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]] }
Expand All @@ -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)

}
)
'{
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down

0 comments on commit df5f46a

Please sign in to comment.