Add discriminator property to generated typeables. Fix tofu-tf#113

leviysoft · Sep 17, 2020 · 45ffd3e · 45ffd3e
1 parent ac65c9e
commit 45ffd3e
Show file tree

Hide file tree

Showing 3 changed files with 61 additions and 13 deletions.
diff --git a/modules/finagle-env/src/test/scala/ru/tinkoff/tschema/finagle/AuthorizationSpec.scala b/modules/finagle-env/src/test/scala/ru/tinkoff/tschema/finagle/AuthorizationSpec.scala
@@ -1,5 +1,7 @@
 package ru.tinkoff.tschema.finagle
 
+import scala.language.reflectiveCalls
+
 import com.twitter.finagle.http.Response
 import monix.eval.Task
 import org.scalatest.flatspec.AnyFlatSpec

diff --git a/modules/swagger/src/main/scala/ru/tinkoff/tschema/swagger/Swagger.scala b/modules/swagger/src/main/scala/ru/tinkoff/tschema/swagger/Swagger.scala
@@ -1,12 +1,12 @@
 package ru.tinkoff.tschema.swagger
 
-import cats.Eval
+import cats.{Endo, Eval}
 import cats.instances.list._
 import cats.syntax.option._
 import derevo.Derivation
 import magnolia.{CaseClass, Magnolia, SealedTrait, TypeName}
+import ru.tinkoff.tschema.swagger._
 import ru.tinkoff.tschema.swagger.SwaggerTypeable.{Config, seq}
-import ru.tinkoff.tschema.swagger.{SwaggerTypeable, _}
 
 object Swagger extends Derivation[SwaggerTypeable] {
   type Typeclass[T] = SwaggerTypeable[T]
@@ -46,9 +46,11 @@ object Swagger extends Derivation[SwaggerTypeable] {
                   typ = Eval.later(param.typeclass.typ)
                 )
               }.toVector,
-              required = Eval.later(caseClass.parameters.toVector.collect {
-                case prop if !prop.typeclass.optional && prop.default.isEmpty => cfg.propMod(prop.label)
-              })
+              required = Eval.later(
+                caseClass.parameters.toVector.collect {
+                  case prop if !prop.typeclass.optional && prop.default.isEmpty => cfg.propMod(prop.label)
+                }
+              )
             )
           )
         )
@@ -59,14 +61,38 @@ object Swagger extends Derivation[SwaggerTypeable] {
       desc: DescribeTypeable[T] = DescribeTypeable.empty[T]
   ): Typeclass[T] =
     new Typeclass[T] {
+      private val typeName = calcTypeName(sealedTrait.typeName, cfg)
+
+      private val dtr = cfg.discriminator.map(cfg.propMod)
+
+      private def addDiscriminator(caseName: String): Endo[SwaggerType] = {
+        case SwaggerObject(properties, required, discriminator) =>
+          SwaggerObject(
+            properties ++ dtr.map(
+              SwaggerProperty(
+                _,
+                None,
+                Eval.later(SwaggerPrimitive.string.mod(_.copy(pattern = Some(cfg.altMod(caseName)))))
+              )
+            ),
+            required.map(_ ++ dtr),
+            discriminator
+          )
+        case SwaggerRef(_, descr, typ)                          =>
+          SwaggerRef(s"$typeName.$caseName", descr, typ.map(addDiscriminator(caseName)))
+        case other                                              => other
+      }
+
       lazy val typ: SwaggerType =
         SwaggerRef(
-          name = calcTypeName(sealedTrait.typeName, cfg),
+          name = typeName,
           descr = desc.whole,
           typ = Eval.now(
             SwaggerOneOf(
               alts = sealedTrait.subtypes.map { sub =>
-                cfg.altMod(sub.typeName.short).some -> Eval.later(sub.typeclass.typ)
+                cfg.altMod(sub.typeName.short).some -> Eval.later(
+                  sub.typeclass.updateTyp(addDiscriminator(sub.typeName.short)).typ
+                )
               }.toVector,
               discriminator = cfg.discriminator
             )

diff --git a/modules/swagger/src/main/scala/ru/tinkoff/tschema/swagger/SwaggerTypeable.scala b/modules/swagger/src/main/scala/ru/tinkoff/tschema/swagger/SwaggerTypeable.scala
@@ -3,7 +3,7 @@ package ru.tinkoff.tschema.swagger
 import java.time.{Instant, LocalDate, LocalDateTime, LocalTime, OffsetDateTime, OffsetTime, ZonedDateTime}
 import java.util.{Date, UUID}
 
-import cats.Eval
+import cats.{Endo, Eval}
 import enumeratum.{Enum, EnumEntry}
 import io.circe.JsonObject
 import cats.syntax.option._
@@ -274,11 +274,16 @@ object GenericSwaggerTypeable {
   implicit def genericProductTypeable[T, L <: HList](implicit
       lgen: LabelledGeneric.Aux[T, L],
       list: HListProps[L],
-      descr: DescribeTypeable[T] = DescribeTypeable.empty[T]
+      descr: DescribeTypeable[T] = DescribeTypeable.empty[T],
+      cfg: Config = SwaggerTypeable.defaultConfig
   ): GenericSwaggerTypeable[T] = {
-    def required = Eval.later(list.props.collect { case (name, tt) if !tt.value._2 => name }.toVector)
+    def required = Eval.later(list.props.collect {
+      case (name, tt) if !tt.value._2 => cfg.propMod(name)
+    }.toVector)
 
-    def props = list.props.map { case (name, tt) => SwaggerProperty(name, descr.element(name), tt.map(_._1)) }.toVector
+    def props = list.props.map {
+      case (name, tt) => SwaggerProperty(name, descr.element(name), tt.map(_._1))
+    }.toVector
 
     make[T](SwaggerObject(props, required), descr.whole)
   }
@@ -301,14 +306,29 @@ object GenericSwaggerTypeable {
       sum: CoproductAlts[C],
       cfg: Config = SwaggerTypeable.defaultConfig,
       descr: DescribeTypeable[T] = DescribeTypeable.empty[T]
-  ): GenericSwaggerTypeable[T] =
+  ): GenericSwaggerTypeable[T] = {
+    val dtr = cfg.discriminator.map(cfg.propMod)
+
+    val addDiscriminator: Endo[SwaggerType] = {
+      case SwaggerObject(properties, required, discriminator) =>
+        SwaggerObject(
+          properties ++ dtr.map(SwaggerProperty(_, None, Eval.later(SwaggerPrimitive.string))),
+          required.map(_ ++ dtr),
+          discriminator
+        )
+      case other                                              => other
+    }
+
     make[T](
       SwaggerOneOf(
-        sum.alts.map { case (name, typ) => name -> typ.map(_.describeWith(name.flatMap(descr.element))) }.toVector,
+        sum.alts.map {
+          case (name, typ) => name -> typ.map(addDiscriminator).map(_.describeWith(name.flatMap(descr.element)))
+        }.toVector,
         cfg.discriminator
       ),
       descr.whole
     )
+  }
 }
 
 sealed trait CirceSwaggerTypeableInstances {