Add discriminator to union-type schema

src/Data/OpenApi/Internal.hs

@@ -670,7 +670,7 @@ data Discriminator = Discriminator
     _discriminatorPropertyName :: Text
 
     -- | An object to hold mappings between payload values and schema names or references.
-  , _discriminatorMapping :: InsOrdHashMap Text Text
+  , _discriminatorMapping :: InsOrdHashMap Text ReferenceToSchema
   } deriving (Eq, Show, Generic, Data, Typeable)
 
 -- | A @'Schema'@ with an optional name.
@@ -947,6 +947,9 @@ instance Hashable ExternalDocs
 newtype Reference = Reference { getReference :: Text }
   deriving (Eq, Show, Data, Typeable)
 
+data ReferenceToSchema = ReferenceToSchema { getReferenceToSchema :: Reference }
+  deriving (Eq, Show, Data, Typeable)
+
 data Referenced a
   = Ref Reference
   | Inline a
@@ -1190,6 +1193,9 @@ instance ToJSON Xml where
 instance ToJSON Discriminator where
   toJSON = genericToJSON (jsonPrefix "Discriminator")
 
+instance ToJSON ReferenceToSchema where
+  toJSON (ReferenceToSchema (Reference t)) = String $ "#/components/schemas/" <> t
+
 instance ToJSON OAuth2ImplicitFlow where
   toJSON = genericToJSON (jsonPrefix "OAuth2ImplicitFlow")
 
@@ -1242,6 +1248,10 @@ instance FromJSON ExternalDocs where
 instance FromJSON Discriminator where
   parseJSON = genericParseJSON (jsonPrefix "Discriminator")
 
+instance FromJSON ReferenceToSchema where
+  parseJSON (String s) | Text.isPrefixOf "#/components/schemas/" s = pure $ ReferenceToSchema . Reference $ Text.drop 21 s
+  parseJSON _ = fail "FromJSON ReferenceToSchema"
+
 instance FromJSON OAuth2ImplicitFlow where
   parseJSON = genericParseJSON (jsonPrefix "OAuth2ImplicitFlow")
 

src/Data/OpenApi/Internal/Schema.hs

@@ -34,6 +34,7 @@ import Data.Aeson (Object (..), SumEncoding (..), ToJSON (..), ToJSONKey (..),
 import Data.Char
 import Data.Data (Data)
 import Data.Foldable (traverse_)
+import Data.Traversable (for)
 import Data.HashMap.Strict (HashMap)
 import qualified Data.HashMap.Strict as HashMap
 import           "unordered-containers" Data.HashSet (HashSet)
@@ -915,8 +916,8 @@ nullarySchema = mempty
   & type_ ?~ OpenApiArray
   & items ?~ OpenApiItemsArray []
 
-gtoNamedSchema :: GToSchema f => SchemaOptions -> Proxy f -> NamedSchema
-gtoNamedSchema opts proxy = undeclare $ gdeclareNamedSchema opts proxy mempty
+gtoNamedSchema :: GToSchema f => SchemaOptions -> Proxy f -> Declare (Definitions Schema) NamedSchema
+gtoNamedSchema opts proxy = gdeclareNamedSchema opts proxy mempty
 
 gdeclareSchema :: GToSchema f => SchemaOptions -> Proxy f -> Declare (Definitions Schema) Schema
 gdeclareSchema opts proxy = _namedSchemaSchema <$> gdeclareNamedSchema opts proxy mempty
@@ -957,7 +958,8 @@ instance (Selector s, GToSchema f, GToSchema (S1 s f)) => GToSchema (C1 c (S1 s
 
 gdeclareSchemaRef :: GToSchema a => SchemaOptions -> Proxy a -> Declare (Definitions Schema) (Referenced Schema)
 gdeclareSchemaRef opts proxy = do
-  case gtoNamedSchema opts proxy of
+  namedSchema <- gtoNamedSchema opts proxy
+  case namedSchema of
     NamedSchema (Just name) schema -> do
       -- This check is very important as it allows generically
       -- derive used definitions for recursive schemas.
@@ -1020,22 +1022,51 @@ instance ( GSumToSchema f
   -- Aeson does not unwrap unary record in sum types.
   gdeclareNamedSchema opts p s = gdeclareNamedSumSchema (opts { unwrapUnaryRecords = False } )p s
 
+-- | Convert inline or ref to ref
+toReferenced :: T.Text -> Referenced Schema -> Declare (Definitions Schema) (Referenced Schema)
+toReferenced _ ref@(Ref _) = pure ref
+toReferenced constructorName r@(Inline schema) = do
+  defs <- look
+  case InsOrdHashMap.lookup constructorName defs of
+    Just schemaAtRef
+      -- Same structure at ref
+      | schemaAtRef == schema -> pure $ Ref $ Reference constructorName
+      -- Same name but structures are different
+      | otherwise -> toReferenced (constructorName <> "_") r -- Modify ref-name with undercore at end
+    Nothing -> do
+      declare $ InsOrdHashMap.insert constructorName schema defs
+      pure $ Ref $ Reference constructorName
+
 gdeclareNamedSumSchema :: GSumToSchema f => SchemaOptions -> Proxy f -> Schema -> Declare (Definitions Schema) NamedSchema
 gdeclareNamedSumSchema opts proxy _
   | allNullaryToStringTag opts && allNullary = pure $ unnamed (toStringTag sumSchemas)
   | otherwise = do
-    (schemas, _) <- runWriterT declareSumSchema
+    (schemas', _) <- runWriterT declareSumSchema
+    schemas <- for schemas' $ \(name, schema) -> do
+            newSchema <- toReferenced name schema
+            pure (name, newSchema)
     return $ unnamed $ mempty
       & type_ ?~ OpenApiObject
       & oneOf ?~ (snd <$> schemas)
+      & discriminator .~ getDiscriminator schemas
   where
     declareSumSchema = gsumToSchema opts proxy
     (sumSchemas, All allNullary) = undeclare (runWriterT declareSumSchema)
-
     toStringTag schemas = mempty
       & type_ ?~ OpenApiString
       & enum_ ?~ map (String . fst) sumSchemas
 
+    tagName = case sumEncoding opts of
+                TaggedObject tagField _ -> Just tagField
+                _ -> Nothing
+    getDiscriminator schemas = do
+          tagPropertyName <- tagName
+          pure Discriminator { _discriminatorPropertyName = T.pack tagPropertyName
+                             , _discriminatorMapping = InsOrdHashMap.fromList 
+                                                        [(name, ReferenceToSchema ref) | (name, Ref ref) <- schemas]
+                             }
+
+
 type AllNullary = All
 
 class GSumToSchema (f :: * -> *)  where

src/Data/OpenApi/Internal/Schema/Validation.hs

@@ -383,9 +383,9 @@ validateObject o = withSchema $ \sch ->
   case sch ^. discriminator of
     Just (Discriminator pname types) -> case fromJSON <$> lookupKey pname o of
       Just (Success pvalue) ->
-        let ref = fromMaybe pvalue $ InsOrdHashMap.lookup pvalue types
+        let ref :: Referenced Schema = maybe (Ref (Reference pname)) (Ref . getReferenceToSchema) (InsOrdHashMap.lookup pvalue types)
         -- TODO ref may be name or reference
-        in validateWithSchemaRef (Ref (Reference ref)) (Object o)
+        in validateWithSchemaRef ref (Object o)
       Just (Error msg)   -> invalid ("failed to parse discriminator property " ++ show pname ++ ": " ++ show msg)
       Nothing            -> invalid ("discriminator property " ++ show pname ++ "is missing")
     Nothing -> do