packages feed

autodocodec-schema-0.2.0.1: src/Autodocodec/Schema.hs

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

module Autodocodec.Schema
  ( jsonSchemaViaCodec,
    jsonSchemaVia,
    jsonObjectSchemaViaCodec,
    jsonObjectSchemaVia,
    JSONSchema (..),
    ObjectSchema (..),
    KeyRequirement (..),
    validateAccordingTo,
    validateObjectAccordingTo,
  )
where

import Autodocodec
import qualified Autodocodec.Aeson.Compat as Compat
import Control.Monad
import Control.Monad.State
import Data.Aeson (FromJSON (..), ToJSON (..))
import qualified Data.Aeson as JSON
import qualified Data.Aeson.Types as JSON
import Data.Foldable
import qualified Data.HashMap.Strict as HM
import Data.List.NonEmpty (NonEmpty (..))
import qualified Data.List.NonEmpty as NE
import Data.Map (Map)
import qualified Data.Map as M
import Data.Maybe
import Data.Scientific
import Data.Set (Set)
import qualified Data.Set as S
import Data.Text (Text)
import qualified Data.Text as T
import Data.Validity
import Data.Validity.Aeson ()
import Data.Validity.Containers ()
import Data.Validity.Text ()
import GHC.Generics (Generic)

-- | A JSON Schema
--
-- http://json-schema.org/understanding-json-schema/reference/index.html
--
-- Contrary to a 'Codec', values of this type should be finite.
--
-- NOTE: This schema roundtrips to JSON, but it cannot expres everything that a fully-featured json-schema may be able to express.
data JSONSchema
  = AnySchema
  | NullSchema
  | BoolSchema
  | StringSchema
  | IntegerSchema !(Bounds Integer)
  | NumberSchema !(Bounds Scientific)
  | ArraySchema !JSONSchema
  | MapSchema !JSONSchema
  | -- | This needs to be a list because keys should stay in their original ordering.
    ObjectSchema !ObjectSchema
  | ValueSchema !JSON.Value
  | AnyOfSchema !(NonEmpty JSONSchema)
  | OneOfSchema !(NonEmpty JSONSchema)
  | CommentSchema !Text !JSONSchema
  | RefSchema !Text
  | WithDefSchema !(Map Text JSONSchema) !JSONSchema
  deriving (Show, Eq, Ord, Generic)

instance Validity JSONSchema where
  validate js =
    mconcat
      [ genericValidate js,
        declare "never has two nested comments" $ case js of
          CommentSchema _ (CommentSchema _ _) -> False
          _ -> True,
        case js of
          AnyOfSchema cs -> declare "there are 2 of more choices" $ length cs >= 2
          OneOfSchema cs -> declare "there are 2 of more choices" $ length cs >= 2
          _ -> valid
      ]

instance ToJSON JSONSchema where
  toJSON = JSON.object . go
    where
      go :: JSONSchema -> [JSON.Pair]
      go = \case
        AnySchema -> []
        NullSchema -> ["type" JSON..= ("null" :: Text)]
        BoolSchema -> ["type" JSON..= ("boolean" :: Text)]
        StringSchema -> ["type" JSON..= ("string" :: Text)]
        IntegerSchema Bounds {..} ->
          catMaybes
            [ Just ("type" JSON..= ("integer" :: Text)),
              ("minimum" JSON..=) <$> boundsLower,
              ("maximum" JSON..=) <$> boundsUpper
            ]
        NumberSchema Bounds {..} ->
          catMaybes
            [ Just ("type" JSON..= ("number" :: Text)),
              ("minimum" JSON..=) <$> boundsLower,
              ("maximum" JSON..=) <$> boundsUpper
            ]
        ArraySchema s ->
          let itemSchemaVal = go s
           in ["type" JSON..= ("array" :: Text), ("items", JSON.object itemSchemaVal)]
        ValueSchema v -> ["const" JSON..= v]
        MapSchema s ->
          let itemSchemaVal = go s
           in ["type" JSON..= ("object" :: Text), "additionalProperties" JSON..= JSON.object itemSchemaVal]
        ObjectSchema os ->
          case toJSON os of
            JSON.Object o -> Compat.toList o
            _ -> [] -- Should not happen.
        AnyOfSchema jcs ->
          let svals :: [JSON.Value]
              svals = map (JSON.object . go) (NE.toList jcs)
              val :: JSON.Value
              val = (JSON.toJSON :: [JSON.Value] -> JSON.Value) svals
           in [("anyOf", val)]
        OneOfSchema jcs ->
          let svals :: [JSON.Value]
              svals = map (JSON.object . go) (NE.toList jcs)
              val :: JSON.Value
              val = (JSON.toJSON :: [JSON.Value] -> JSON.Value) svals
           in [("oneOf", val)]
        (CommentSchema outerComment (CommentSchema innerComment s)) ->
          go (CommentSchema (outerComment <> "\n" <> innerComment) s)
        CommentSchema comment s -> ("$comment" JSON..= comment) : go s
        RefSchema name -> ["$ref" JSON..= (defsPrefix <> name :: Text)]
        WithDefSchema defs s -> ("$defs" JSON..= defs) : go s

instance FromJSON JSONSchema where
  parseJSON = JSON.withObject "JSONSchema" $ \o -> do
    mt <- o JSON..:? "type"
    mc <- o JSON..:? "$comment"
    let commentFunc = maybe id CommentSchema mc
    mdefs <- o JSON..:? "$defs"
    let defsFunc = maybe id WithDefSchema mdefs
    fmap (commentFunc . defsFunc) $ case mt :: Maybe Text of
      Just "null" -> pure NullSchema
      Just "boolean" -> pure BoolSchema
      Just "string" -> pure StringSchema
      Just "integer" -> do
        boundsLower <- o JSON..:? "minimum"
        boundsUpper <- o JSON..:? "maximum"
        pure $ IntegerSchema Bounds {..}
      Just "number" -> do
        boundsLower <- o JSON..:? "minimum"
        boundsUpper <- o JSON..:? "maximum"
        pure $ NumberSchema Bounds {..}
      Just "array" -> do
        mI <- o JSON..:? "items"
        case mI of
          Nothing -> pure $ ArraySchema AnySchema
          Just is -> pure $ ArraySchema is
      Just "object" -> do
        mAdditional <- o JSON..:? "additionalProperties"
        case mAdditional of
          Nothing -> ObjectSchema <$> parseJSON (JSON.Object o)
          Just additional -> pure $ MapSchema additional
      Nothing -> do
        mAny <- o JSON..:? "anyOf"
        case mAny of
          Just anies -> pure $ AnyOfSchema anies
          Nothing -> do
            mOne <- o JSON..:? "oneOf"
            case mOne of
              Just ones -> pure $ OneOfSchema ones
              Nothing -> do
                let mConst = Compat.lookupKey "const" o
                case mConst of
                  Just constant -> pure $ ValueSchema constant
                  Nothing -> do
                    mRef <- o JSON..:? "$ref"
                    pure $ case mRef of
                      Just ref -> case T.stripPrefix defsPrefix ref of
                        Just name -> RefSchema name
                        Nothing -> AnySchema
                      Nothing -> AnySchema
      t -> fail $ "unknown schema type: " <> show t

data ObjectSchema
  = ObjectKeySchema !Text !KeyRequirement !JSONSchema !(Maybe Text)
  | ObjectAnySchema -- For 'pure'
  | ObjectAnyOfSchema !(NonEmpty ObjectSchema)
  | ObjectOneOfSchema !(NonEmpty ObjectSchema)
  | ObjectAllOfSchema !(NonEmpty ObjectSchema)
  deriving (Show, Eq, Ord, Generic)

instance Validity ObjectSchema

instance FromJSON ObjectSchema where
  parseJSON = JSON.withObject "ObjectSchema" go
    where
      go :: JSON.Object -> JSON.Parser ObjectSchema
      go o = do
        t <- o JSON..: "type"
        guard $ t == ("object" :: Text)
        mAllOf <- o JSON..:? "allOf"
        case mAllOf of
          Just ao -> do
            ne <- parseJSON ao
            ObjectAllOfSchema <$> mapM go ne
          Nothing -> do
            mAnyOf <- o JSON..:? "anyOf"
            case mAnyOf of
              Just anies -> do
                ne <- parseJSON anies
                ObjectAnyOfSchema <$> mapM go ne
              Nothing -> do
                mOneOf <- o JSON..:? "oneOf"
                case mOneOf of
                  Just ones -> do
                    ne <- parseJSON ones
                    ObjectOneOfSchema <$> mapM go ne
                  Nothing -> do
                    props <- o JSON..:? "properties" JSON..!= HM.empty
                    reqs <- o JSON..:? "required" JSON..!= []
                    let keySchemaFor k v = do
                          ks <- parseJSON v
                          let (mDoc, ks') = case ks of
                                CommentSchema doc ks'' -> (Just doc, ks'')
                                _ -> (Nothing, ks)
                          pure $
                            if k `elem` reqs
                              then ObjectKeySchema k Required ks' mDoc
                              else ObjectKeySchema k (Optional Nothing) ks' mDoc
                    keySchemas <- mapM (uncurry keySchemaFor) (HM.toList props)
                    pure $ case NE.nonEmpty keySchemas of
                      Nothing -> ObjectAnySchema
                      Just (el :| []) -> el
                      Just ne -> ObjectAllOfSchema ne

instance ToJSON ObjectSchema where
  toJSON = JSON.object . (("type" JSON..= ("object" :: Text)) :) . go
    where
      go :: ObjectSchema -> [JSON.Pair]
      go = \case
        ObjectAnySchema -> []
        ObjectKeySchema k kr ks mDoc ->
          let (propVal, req) = keySchemaToPieces (k, kr, ks, mDoc)
           in -- TODO deal with the default value somehow.
              concat [["properties" JSON..= JSON.object [Compat.toKey k JSON..= propVal]], ["required" JSON..= [k] | req]]
        ObjectAnyOfSchema ne -> ["anyOf" JSON..= NE.map toJSON ne]
        ObjectOneOfSchema ne -> ["oneOf" JSON..= NE.map toJSON ne]
        ObjectAllOfSchema ne ->
          case mapM parseAndObjectKeySchema (NE.toList ne) of
            Nothing -> ["allOf" JSON..= NE.map toJSON ne]
            Just ne' ->
              let f (hm, l) tup@(k, _, _, _) =
                    let (propVal, req) = keySchemaToPieces tup
                     in (HM.insert k propVal hm, if req then k : l else l)
                  (propValMap, reqs) = foldl' f (HM.empty, []) (concat ne')
               in concat [["properties" JSON..= propValMap], ["required" JSON..= reqs | not $ null reqs]]

      keySchemaToPieces :: (Text, KeyRequirement, JSONSchema, Maybe Text) -> (JSON.Value, Bool)
      keySchemaToPieces (_, kr, ks, mDoc) =
        let propVal = toJSON (maybe id CommentSchema mDoc ks)
         in (propVal, kr == Required)

      parseAndObjectKeySchema :: ObjectSchema -> Maybe [(Text, KeyRequirement, JSONSchema, Maybe Text)]
      parseAndObjectKeySchema = \case
        ObjectKeySchema k kr ks mDoc -> Just [(k, kr, ks, mDoc)]
        ObjectAllOfSchema os -> concat <$> mapM parseAndObjectKeySchema os
        _ -> Nothing

defsPrefix :: Text
defsPrefix = "#/$defs/"

data KeyRequirement
  = Required
  | Optional !(Maybe JSON.Value) -- Default value
  deriving (Show, Eq, Ord, Generic)

instance Validity KeyRequirement

jsonSchemaViaCodec :: forall a. (HasCodec a) => JSONSchema
jsonSchemaViaCodec = jsonSchemaVia (codec @a)

jsonSchemaVia :: ValueCodec input output -> JSONSchema
jsonSchemaVia = (`evalState` S.empty) . goValue

jsonObjectSchemaViaCodec :: forall a. (HasObjectCodec a) => ObjectSchema
jsonObjectSchemaViaCodec = jsonObjectSchemaVia (objectCodec @a)

jsonObjectSchemaVia :: ObjectCodec input output -> ObjectSchema
jsonObjectSchemaVia = (`evalState` S.empty) . goObject

goValue :: ValueCodec input output -> State (Set Text) JSONSchema
goValue = \case
  NullCodec -> pure NullSchema
  BoolCodec mname -> pure $ maybe id CommentSchema mname BoolSchema
  StringCodec mname -> pure $ maybe id CommentSchema mname StringSchema
  IntegerCodec mname mBounds -> pure $ maybe id CommentSchema mname $ IntegerSchema mBounds
  NumberCodec mname mBounds -> pure $ maybe id CommentSchema mname $ NumberSchema mBounds
  ArrayOfCodec mname c -> do
    s <- goValue c
    pure $ maybe id CommentSchema mname $ ArraySchema s
  ObjectOfCodec mname oc -> do
    s <- goObject oc
    pure $ maybe id CommentSchema mname $ ObjectSchema s
  HashMapCodec c -> MapSchema <$> goValue c
  MapCodec c -> MapSchema <$> goValue c
  ValueCodec -> pure AnySchema
  EqCodec value c -> pure $ ValueSchema (toJSONVia c value)
  EitherCodec u c1 c2 -> do
    s1 <- goValue c1
    s2 <- goValue c2
    pure $ case u of
      DisjointUnion -> OneOfSchema (goOneOf (s1 :| [s2]))
      PossiblyJointUnion -> AnyOfSchema (goAnyOf (s1 :| [s2]))
  BimapCodec _ _ c -> goValue c
  CommentCodec t c -> CommentSchema t <$> goValue c
  ReferenceCodec name c -> do
    alreadySeen <- gets (S.member name)
    if alreadySeen
      then pure $ RefSchema name
      else do
        modify (S.insert name)
        s <- goValue c
        pure $ WithDefSchema (M.singleton name s) (RefSchema name)
  where
    goAnyOf :: NonEmpty JSONSchema -> NonEmpty JSONSchema
    goAnyOf (s :| rest) = case NE.nonEmpty rest of
      Nothing -> goSingle s
      Just ne -> goSingle s <> goAnyOf ne
      where
        goSingle :: JSONSchema -> NonEmpty JSONSchema
        goSingle = \case
          AnyOfSchema ss -> goAnyOf ss
          s' -> s' :| []
    goOneOf :: NonEmpty JSONSchema -> NonEmpty JSONSchema
    goOneOf (s :| rest) = case NE.nonEmpty rest of
      Nothing -> goSingle s
      Just ne -> goSingle s <> goOneOf ne
      where
        goSingle :: JSONSchema -> NonEmpty JSONSchema
        goSingle = \case
          OneOfSchema ss -> goOneOf ss
          s' -> s' :| []

goObject :: ObjectCodec input output -> State (Set Text) ObjectSchema
goObject = \case
  RequiredKeyCodec k c mdoc -> do
    s <- goValue c
    pure $ ObjectKeySchema k Required s mdoc
  OptionalKeyCodec k c mdoc -> do
    s <- goValue c
    pure $ ObjectKeySchema k (Optional Nothing) s mdoc
  OptionalKeyWithDefaultCodec k c mr mdoc -> do
    s <- goValue c
    pure $ ObjectKeySchema k (Optional (Just (toJSONVia c mr))) s mdoc
  OptionalKeyWithOmittedDefaultCodec k c defaultValue mDoc -> goObject (optionalKeyWithDefaultCodec k c defaultValue mDoc)
  BimapCodec _ _ c -> goObject c
  EitherCodec u oc1 oc2 -> do
    os1 <- goObject oc1
    os2 <- goObject oc2
    pure $ case u of
      DisjointUnion -> ObjectOneOfSchema (goObjectOneOf (os1 :| [os2]))
      PossiblyJointUnion -> ObjectAnyOfSchema (goObjectAnyOf (os1 :| [os2]))
  DiscriminatedUnionCodec pn _ m -> do
    let mkSchema dName (_, oc) =
          goObject $ oc *> (requiredFieldWith' pn (literalTextCodec dName) .= const dName)
    ss <- HM.traverseWithKey mkSchema m
    pure $ case NE.nonEmpty $ toList ss of
      Nothing -> ObjectAnySchema
      Just ss' -> ObjectOneOfSchema $ goObjectOneOf ss'
  PureCodec _ -> pure ObjectAnySchema
  ApCodec oc1 oc2 -> do
    os1 <- goObject oc1
    os2 <- goObject oc2
    pure $ ObjectAllOfSchema (goObjectAllOf (os1 :| [os2]))
  where
    goObjectAnyOf :: NonEmpty ObjectSchema -> NonEmpty ObjectSchema
    goObjectAnyOf (s :| rest) = case NE.nonEmpty rest of
      Nothing -> goSingle s
      Just ne -> goSingle s <> goObjectAnyOf ne
      where
        goSingle :: ObjectSchema -> NonEmpty ObjectSchema
        goSingle = \case
          ObjectAnyOfSchema ss -> goObjectAnyOf ss
          s' -> s' :| []

    goObjectOneOf :: NonEmpty ObjectSchema -> NonEmpty ObjectSchema
    goObjectOneOf (s :| rest) = case NE.nonEmpty rest of
      Nothing -> goSingle s
      Just ne -> goSingle s <> goObjectOneOf ne
      where
        goSingle :: ObjectSchema -> NonEmpty ObjectSchema
        goSingle = \case
          ObjectOneOfSchema ss -> goObjectOneOf ss
          s' -> s' :| []

    goObjectAllOf :: NonEmpty ObjectSchema -> NonEmpty ObjectSchema
    goObjectAllOf (s :| rest) = case NE.nonEmpty rest of
      Nothing -> goSingle s
      Just ne -> goSingle s <> goObjectAllOf ne
      where
        goSingle :: ObjectSchema -> NonEmpty ObjectSchema
        goSingle = \case
          ObjectAllOfSchema ss -> goObjectAllOf ss
          s' -> s' :| []

validateAccordingTo :: JSON.Value -> JSONSchema -> Bool
validateAccordingTo val schema = (`evalState` M.empty) $ validateValue val schema

validateObjectAccordingTo :: JSON.Value -> JSONSchema -> Bool
validateObjectAccordingTo val schema = (`evalState` M.empty) $ validateValue val schema

validateValue :: JSON.Value -> JSONSchema -> State (Map Text JSONSchema) Bool
validateValue value = \case
  AnySchema -> pure True
  NullSchema -> pure $ value == JSON.Null
  BoolSchema -> pure $ case value of
    JSON.Bool _ -> True
    _ -> False
  StringSchema -> pure $ case value of
    JSON.String _ -> True
    _ -> False
  IntegerSchema bounds -> pure $ case value of
    JSON.Number s -> case checkBounds (fromInteger <$> bounds) s of
      Left _ -> False
      Right _ -> True
    _ -> False
  NumberSchema bounds -> pure $ case value of
    JSON.Number s -> case checkBounds bounds s of
      Left _ -> False
      Right _ -> True
    _ -> False
  ArraySchema as -> case value of
    JSON.Array v -> and <$> mapM (`validateValue` as) v
    _ -> pure False
  MapSchema vs -> case value of
    JSON.Object hm -> and <$> mapM (`validateValue` vs) hm
    _ -> pure False
  ObjectSchema os -> case value of
    JSON.Object obj -> validateObject obj os
    _ -> pure False
  ValueSchema v -> pure $ v == value
  AnyOfSchema ss -> or <$> mapM (validateValue value) ss
  OneOfSchema ss -> (== 1) . length . NE.filter id <$> mapM (validateValue value) ss
  CommentSchema _ s -> validateValue value s
  RefSchema name -> do
    mSchema <- gets (M.lookup name)
    case mSchema of
      Nothing -> pure False -- Referred to a schema that's not defined, we have no choice but to reject the value.
      Just s -> validateValue value s
  WithDefSchema defs s -> do
    modify (M.union defs)
    validateValue value s

validateObject :: JSON.Object -> ObjectSchema -> State (Map Text JSONSchema) Bool
validateObject obj = \case
  ObjectAnySchema -> pure True
  ObjectKeySchema key kr ks _ -> case Compat.lookupKey (Compat.toKey key) obj of
    Nothing -> case kr of
      Required -> pure False
      Optional _ -> pure True
    Just value' -> validateValue value' ks
  ObjectAllOfSchema ne -> and <$> mapM (validateObject obj) ne
  ObjectAnyOfSchema ne -> or <$> mapM (validateObject obj) ne
  ObjectOneOfSchema ne -> (== 1) . length . NE.filter id <$> mapM (validateObject obj) ne