added simpleEnum annotation (#564)

* added simpleEnum annotation * added automaticallyAdded * added automaticallyAdded * added tests * fmt * restored recordName --------- Co-authored-by: Daniel Vigovszky <[email protected]>
zio · Sep 5, 2023 · bc2c20e · bc2c20e
1 parent 1861252
commit bc2c20e
Show file tree

Hide file tree

Showing 8 changed files with 109 additions and 36 deletions.
diff --git a/zio-schema-derivation/shared/src/main/scala-2/zio/schema/DeriveSchema.scala b/zio-schema-derivation/shared/src/main/scala-2/zio/schema/DeriveSchema.scala
@@ -501,23 +501,53 @@ object DeriveSchema {
         }
       }
 
+      val isSimpleEnum: Boolean =
+        !tpe.typeSymbol.asClass.knownDirectSubclasses.map { subtype =>
+          subtype.typeSignature.decls.sorted.collect {
+            case p: TermSymbol if p.isCaseAccessor && !p.isMethod => p
+          }.size
+        }.exists(_ > 0)
+
+      val hasSimpleEnum: Boolean =
+        tpe.typeSymbol.annotations.exists(_.tree.tpe =:= typeOf[_root_.zio.schema.annotation.simpleEnum])
+
       @nowarn
-      val typeAnnotations: List[Tree] =
-        tpe.typeSymbol.annotations.collect {
-          case annotation if !(annotation.tree.tpe <:< JavaAnnotationTpe) =>
-            annotation.tree match {
-              case q"new $annConstructor(..$annotationArgs)" =>
-                q"new ${annConstructor.tpe.typeSymbol}(..$annotationArgs)"
-              case q"new $annConstructor()" =>
-                q"new ${annConstructor.tpe.typeSymbol}()"
-              case tree =>
-                c.warning(c.enclosingPosition, s"Unhandled annotation tree $tree")
-                EmptyTree
-            }
-          case annotation =>
-            c.warning(c.enclosingPosition, s"Unhandled annotation ${annotation.tree}")
-            EmptyTree
-        }.filter(_ != EmptyTree)
+      val typeAnnotations: List[Tree] = (isSimpleEnum, hasSimpleEnum) match {
+        case (true, false) =>
+          tpe.typeSymbol.annotations.collect {
+            case annotation if !(annotation.tree.tpe <:< JavaAnnotationTpe) =>
+              annotation.tree match {
+                case q"new $annConstructor(..$annotationArgs)" =>
+                  q"new ${annConstructor.tpe.typeSymbol}(..$annotationArgs)"
+                case q"new $annConstructor()" =>
+                  q"new ${annConstructor.tpe.typeSymbol}()"
+                case tree =>
+                  c.warning(c.enclosingPosition, s"Unhandled annotation tree $tree")
+                  EmptyTree
+              }
+            case annotation =>
+              c.warning(c.enclosingPosition, s"Unhandled annotation ${annotation.tree}")
+              EmptyTree
+          }.filter(_ != EmptyTree).+:(q"new _root_.zio.schema.annotation.simpleEnum(true)")
+        case (false, true) =>
+          c.abort(c.enclosingPosition, s"${show(tpe)} must be a simple Enum")
+        case _ =>
+          tpe.typeSymbol.annotations.collect {
+            case annotation if !(annotation.tree.tpe <:< JavaAnnotationTpe) =>
+              annotation.tree match {
+                case q"new $annConstructor(..$annotationArgs)" =>
+                  q"new ${annConstructor.tpe.typeSymbol}(..$annotationArgs)"
+                case q"new $annConstructor()" =>
+                  q"new ${annConstructor.tpe.typeSymbol}()"
+                case tree =>
+                  c.warning(c.enclosingPosition, s"Unhandled annotation tree $tree")
+                  EmptyTree
+              }
+            case annotation =>
+              c.warning(c.enclosingPosition, s"Unhandled annotation ${annotation.tree}")
+              EmptyTree
+          }.filter(_ != EmptyTree)
+      }
 
       val selfRefName  = c.freshName("ref")
       val selfRefIdent = Ident(TermName(selfRefName))

diff --git a/zio-schema-derivation/shared/src/main/scala-3/zio/schema/DeriveSchema.scala b/zio-schema-derivation/shared/src/main/scala-3/zio/schema/DeriveSchema.scala
@@ -293,7 +293,14 @@ private case class DeriveSchema()(using val ctx: Quotes) extends ReflectionUtils
 
     val cases = typesAndLabels.map { case (tpe, label) => deriveCase[T](tpe, label, newStack) }
 
-    val annotationExprs = TypeRepr.of[T].typeSymbol.annotations.filter(filterAnnotation).map(_.asExpr)
+    val isSimpleEnum: Boolean = !TypeRepr.of[T].typeSymbol.children.map(_.declaredFields.length).exists( _ > 0 )
+    val hasSimpleEnumAnn: Boolean = TypeRepr.of[T].typeSymbol.hasAnnotation(TypeRepr.of[_root_.zio.schema.annotation.simpleEnum].typeSymbol)
+
+    val annotationExprs = (isSimpleEnum, hasSimpleEnumAnn) match {
+      case (true, false) => TypeRepr.of[T].typeSymbol.annotations.filter(filterAnnotation).map(_.asExpr).+:('{_root_.zio.schema.annotation.simpleEnum(true)})
+      case (false, true) => throw new Exception(s"${TypeRepr.of[T].typeSymbol.name} must be a simple Enum")
+      case _             => TypeRepr.of[T].typeSymbol.annotations.filter(filterAnnotation).map(_.asExpr)
+    }
     val annotations = '{ zio.Chunk.fromIterable(${Expr.ofSeq(annotationExprs)}) }
 
     val typeInfo = '{TypeId.parse(${Expr(TypeRepr.of[T].show)})}

diff --git a/zio-schema-derivation/shared/src/test/scala/zio/schema/DeriveSchemaSpec.scala b/zio-schema-derivation/shared/src/test/scala/zio/schema/DeriveSchemaSpec.scala
@@ -3,7 +3,7 @@ package zio.schema
 import scala.annotation.Annotation
 
 import zio.Chunk
-import zio.schema.annotation.{ fieldName, optionalField }
+import zio.schema.annotation.{ fieldName, optionalField, simpleEnum }
 import zio.test._
 
 object DeriveSchemaSpec extends ZIOSpecDefault with VersionSpecificDeriveSchemaSpec {
@@ -243,6 +243,13 @@ object DeriveSchemaSpec extends ZIOSpecDefault with VersionSpecificDeriveSchemaS
     implicit val schema: Schema[OptionalField] = DeriveSchema.gen[OptionalField]
   }
 
+  @simpleEnum
+  sealed trait SimpleEnum1
+  case class SimpleClass1() extends SimpleEnum1
+
+  sealed trait SimpleEnum2
+  case class SimpleClass2() extends SimpleEnum2
+
   override def spec: Spec[Environment, Any] = suite("DeriveSchemaSpec")(
     suite("Derivation")(
       test("correctly derives case class 0") {
@@ -449,6 +456,14 @@ object DeriveSchemaSpec extends ZIOSpecDefault with VersionSpecificDeriveSchemaS
           )
         }
         assert(derived)(hasSameSchema(expected))
+      },
+      test("correctly derives simpleEnum with annotation") {
+        val derived = DeriveSchema.gen[SimpleEnum1]
+        assertTrue(derived.annotations == Chunk(simpleEnum(false)))
+      },
+      test("correctly derives simpleEnum without annotation") {
+        val derived = DeriveSchema.gen[SimpleEnum2]
+        assertTrue(derived.annotations == Chunk(simpleEnum(true)))
       }
     ),
     versionSpecificSuite

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
@@ -380,7 +380,7 @@ object JsonCodec {
 
     private def enumEncoder[Z](parentSchema: Schema.Enum[Z], cases: Schema.Case[Z, _]*): ZJsonEncoder[Z] =
       // if all cases are CaseClass0, encode as a String
-      if (cases.forall(_.schema.isInstanceOf[Schema.CaseClass0[_]])) {
+      if (parentSchema.annotations.exists(_.isInstanceOf[simpleEnum])) {
         val caseMap: Map[Z, String] = cases
           .filterNot(_.annotations.exists(_.isInstanceOf[transientCase]))
           .map(

diff --git a/zio-schema/shared/src/main/scala/zio/schema/Schema.scala b/zio-schema/shared/src/main/scala/zio/schema/Schema.scala
@@ -161,8 +161,8 @@ object Schema extends SchemaEquality {
 
   def defer[A](schema: => Schema[A]): Schema[A] = Lazy(() => schema)
 
-  def enumeration[A, C <: CaseSet.Aux[A]](id: TypeId, caseSet: C): Schema[A] =
-    EnumN(id, caseSet, Chunk.empty)
+  def enumeration[A, C <: CaseSet.Aux[A]](id: TypeId, caseSet: C, annotations: Chunk[Any] = Chunk.empty): Schema[A] =
+    EnumN(id, caseSet, annotations)
 
   def fail[A](message: String): Schema[A] = Fail(message)
 

diff --git a/zio-schema/shared/src/main/scala/zio/schema/annotation/recordName.scala b/zio-schema/shared/src/main/scala/zio/schema/annotation/recordName.scala
@@ -1,3 +1,5 @@
 package zio.schema.annotation
 
-final case class recordName(name: String) extends scala.annotation.StaticAnnotation
+import scala.annotation.StaticAnnotation
+
+final case class recordName(name: String) extends StaticAnnotation
diff --git a/zio-schema/shared/src/main/scala/zio/schema/annotation/simpleEnum.scala b/zio-schema/shared/src/main/scala/zio/schema/annotation/simpleEnum.scala
@@ -0,0 +1,10 @@
+package zio.schema.annotation
+
+import scala.annotation.StaticAnnotation
+
+/*
+ * Automatically added in sealed traits with only case objects or case class without parameters.
+ * Gives error if it used in other types of enumerations.
+ */
+
+final case class simpleEnum(automaticallyAdded: Boolean = false) extends StaticAnnotation
diff --git a/zio-schema/shared/src/main/scala/zio/schema/meta/ExtensibleMetaSchema.scala b/zio-schema/shared/src/main/scala/zio/schema/meta/ExtensibleMetaSchema.scala
@@ -7,6 +7,7 @@ import scala.collection.mutable
 import zio.constraintless.TypeList
 import zio.prelude._
 import zio.schema._
+import zio.schema.annotation.simpleEnum
 import zio.{ Chunk, ChunkBuilder }
 
 sealed trait ExtensibleMetaSchema[BuiltIn <: TypeList] { self =>
@@ -636,23 +637,31 @@ object ExtensibleMetaSchema {
           materialize(left, refs),
           materialize(right, refs)
         )
-      case ExtensibleMetaSchema.Sum(id, _, elems, _) =>
+      case ExtensibleMetaSchema.Sum(id, _, elems, _) => {
+        val (cases, isSimple) = elems.foldRight[(CaseSet.Aux[Any], Boolean)]((CaseSet.Empty[Any](), true)) {
+          case (Labelled(label, ast), (acc, wasSimple)) => {
+            val _case: Schema.Case[Any, Any] = Schema
+              .Case[Any, Any](
+                label,
+                materialize(ast, refs).asInstanceOf[Schema[Any]],
+                identity[Any],
+                identity[Any],
+                _.isInstanceOf[Any],
+                Chunk.empty
+              )
+            val isSimple: Boolean = _case.schema match {
+              case _: Schema.CaseClass0[_] => true
+              case _                       => false
+            }
+            (CaseSet.Cons(_case, acc), wasSimple && isSimple)
+          }
+        }
         Schema.enumeration[Any, CaseSet.Aux[Any]](
           id,
-          elems.foldRight[CaseSet.Aux[Any]](CaseSet.Empty[Any]()) {
-            case (Labelled(label, ast), acc) =>
-              val _case: Schema.Case[Any, Any] = Schema
-                .Case[Any, Any](
-                  label,
-                  materialize(ast, refs).asInstanceOf[Schema[Any]],
-                  identity[Any],
-                  identity[Any],
-                  _.isInstanceOf[Any],
-                  Chunk.empty
-                )
-              CaseSet.Cons(_case, acc)
-          }
+          cases,
+          if (isSimple) Chunk(new simpleEnum(true)) else Chunk.empty
         )
+      }
       case ExtensibleMetaSchema.Either(_, left, right, _) =>
         Schema.either(
           materialize(left, refs),