From 6b568b3d6c26ee44fb3801ab97f1167cad368f4f Mon Sep 17 00:00:00 2001
From: Daniel Slapman <danslapman@gmail.com>
Date: Wed, 9 Sep 2020 22:08:42 +0300
Subject: [PATCH] Add discriminator property to generated typeables. Fix #113

---
 .../tschema/finagle/AuthorizationSpec.scala   |  2 +
 .../ru/tinkoff/tschema/swagger/Swagger.scala  | 40 ++++++++++++++-----
 .../tschema/swagger/SwaggerTypeable.scala     | 27 +++++++++++--
 3 files changed, 55 insertions(+), 14 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
index 4603adba..ddac3cfa 100644
--- 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
index fb0f8ac6..d02b9a40 100644
--- a/modules/swagger/src/main/scala/ru/tinkoff/tschema/swagger/Swagger.scala
+++ b/modules/swagger/src/main/scala/ru/tinkoff/tschema/swagger/Swagger.scala
@@ -5,8 +5,8 @@ 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]
@@ -33,22 +33,42 @@ object Swagger extends Derivation[SwaggerTypeable] {
       desc: DescribeTypeable[T] = DescribeTypeable.empty[T]
   ): Typeclass[T] =
     new Typeclass[T] {
+      private val dtr = cfg.discriminator.map(cfg.propMod)
+
+      private def addDiscriminatorProperty(props: Vector[SwaggerProperty]): Vector[SwaggerProperty] =
+        dtr.fold(props) { dName =>
+          if (props.exists(_.name == dName))
+            props
+          else props :+ SwaggerProperty(dName, None, Eval.now(SwaggerPrimitive.string))
+        }
+
+      private def discriminatorRequired(required: Vector[String]): Vector[String] =
+        dtr.fold(required) { dName =>
+          if (required.contains(dName)) required else required :+ dName
+        }
+
       lazy val typ: SwaggerType =
         SwaggerRef(
           name = calcTypeName(caseClass.typeName, cfg),
           descr = desc.whole,
           typ = Eval.now(
             SwaggerObject(
-              properties = caseClass.parameters.map { param =>
-                SwaggerProperty(
-                  name = cfg.propMod(param.label),
-                  description = desc.element(param.label),
-                  typ = Eval.later(param.typeclass.typ)
+              properties = addDiscriminatorProperty(
+                caseClass.parameters.map { param =>
+                  SwaggerProperty(
+                    name = cfg.propMod(param.label),
+                    description = desc.element(param.label),
+                    typ = Eval.later(param.typeclass.typ)
+                  )
+                }.toVector
+              ),
+              required = Eval.later(
+                discriminatorRequired(
+                  caseClass.parameters.toVector.collect {
+                    case prop if !prop.typeclass.optional => cfg.propMod(prop.label)
+                  }
                 )
-              }.toVector,
-              required = Eval.later(caseClass.parameters.toVector.collect {
-                case prop if !prop.typeclass.optional => cfg.propMod(prop.label)
-              })
+              )
             )
           )
         )
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
index 4fa2f93f..b97633ef 100644
--- a/modules/swagger/src/main/scala/ru/tinkoff/tschema/swagger/SwaggerTypeable.scala
+++ b/modules/swagger/src/main/scala/ru/tinkoff/tschema/swagger/SwaggerTypeable.scala
@@ -274,11 +274,30 @@ 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 props = list.props.map { case (name, tt) => SwaggerProperty(name, descr.element(name), tt.map(_._1)) }.toVector
+    val dtr = cfg.discriminator.map(cfg.propMod)
+
+    def addDiscriminatorProperty(props: Vector[SwaggerProperty]): Vector[SwaggerProperty] =
+      dtr.fold(props) { dName =>
+        if (props.exists(_.name == dName))
+          props
+        else props :+ SwaggerProperty(dName, None, Eval.now(SwaggerPrimitive.string))
+      }
+
+    def discriminatorRequired(required: Vector[String]): Vector[String] =
+      dtr.fold(required) { dName =>
+        if (required.contains(dName)) required else required :+ dName
+      }
+
+    def required = Eval.later(discriminatorRequired(list.props.collect {
+      case (name, tt) if !tt.value._2 => cfg.propMod(name)
+    }.toVector))
+
+    def props = addDiscriminatorProperty(list.props.map {
+      case (name, tt) => SwaggerProperty(name, descr.element(name), tt.map(_._1))
+    }.toVector)
 
     make[T](SwaggerObject(props, required), descr.whole)
   }