Add schemaGenWithFormats to support custom string generators

Author: cydparser

src/Data/OpenApi/Schema/Generator.hs

@@ -8,7 +8,6 @@ import           Prelude                                 ()
 import           Prelude.Compat
 
 import           Control.Lens.Operators
-import           Control.Monad                           (filterM)
 import           Data.Aeson
 import           Data.Aeson.Types
 import qualified Data.HashMap.Strict.InsOrd              as M
@@ -19,24 +18,29 @@ import qualified Data.Set                                as S
 import           Data.OpenApi
 import           Data.OpenApi.Declare
 import           Data.OpenApi.Internal.Schema.Validation (inferSchemaTypes)
+import           Data.Text                               (Text)
 import qualified Data.Text                               as T
 import qualified Data.Vector                             as V
+import           GHC.Stack                               (HasCallStack)
 import           Test.QuickCheck                         (arbitrary)
 import           Test.QuickCheck.Gen
 import           Test.QuickCheck.Property
 
 -- | Note: 'schemaGen' may 'error', if schema type is not specified,
 -- and cannot be inferred.
-schemaGen :: Definitions Schema -> Schema -> Gen Value
-schemaGen _ schema
+schemaGen :: HasCallStack => Definitions Schema -> Schema -> Gen Value
+schemaGen = schemaGenWithFormats (const Nothing)
+
+schemaGenWithFormats :: HasCallStack => (Format -> Maybe (Gen Text)) -> Definitions Schema -> Schema -> Gen Value
+schemaGenWithFormats _ _ schema
     | Just cases <- schema  ^. enum_  = elements cases
-schemaGen defns schema
+schemaGenWithFormats _ defns schema
     | Just variants <- schema ^. oneOf = schemaGen defns =<< elements (dereference defns <$> variants)
-schemaGen defns schema =
+schemaGenWithFormats formatGen defns schema =
     case schema ^. type_ of
       Nothing ->
         case inferSchemaTypes schema of
-          [ inferredType ] -> schemaGen defns (schema & type_ ?~ inferredType)
+          [ inferredType ] -> schemaGenWithFormats formatGen defns (schema & type_ ?~ inferredType)
           -- Gen is not MonadFail
           _ -> error "unable to infer schema type"
       Just OpenApiBoolean -> Bool <$> elements [True, False]
@@ -63,12 +67,16 @@ schemaGen defns schema =
                       minLength' = fromMaybe 0 $ fromInteger <$> schema ^. minItems
                       maxLength' = fromMaybe size $ fromInteger <$> schema ^. maxItems
                   arrayLength <- choose (minLength', max minLength' maxLength')
-                  generatedArray <- vectorOf arrayLength $ schemaGen defns itemSchema
+                  generatedArray <- vectorOf arrayLength $ schemaGenWithFormats formatGen defns itemSchema
                   return . Array $ V.fromList generatedArray
               OpenApiItemsArray refs ->
-                  let itemGens = schemaGen defns . dereference defns <$> refs
+                  let itemGens = schemaGenWithFormats formatGen defns . dereference defns <$> refs
                   in fmap (Array . V.fromList) $ sequence itemGens
-      Just OpenApiString -> do
+        | otherwise -> error "invalid array"
+      Just OpenApiString
+        | Just gen <- formatGen =<< schema ^. format ->
+            String <$> gen
+        | otherwise -> do
         size <- getSize
         let minLength' = fromMaybe 0 $ fromInteger <$> schema ^. minLength
         let maxLength' = fromMaybe size $ fromInteger <$> schema ^. maxLength
@@ -88,11 +96,11 @@ schemaGen defns schema =
           numProps <- choose (minProps', max minProps' maxProps')
           let presentKeys = take numProps $ S.toList reqKeys ++ shuffledOptional
           let presentProps = M.filterWithKey (\k _ -> k `elem` presentKeys) props
-          let gens = schemaGen defns <$> presentProps
+          let gens = schemaGenWithFormats formatGen defns <$> presentProps
           additionalGens <- case schema ^. additionalProperties of
             Just (AdditionalPropertiesSchema addlSchema) -> do
               additionalKeys <- sequence . take (numProps - length presentProps) . repeat $ T.pack <$> arbitrary
-              return . M.fromList $ zip additionalKeys (repeat . schemaGen defns $ dereference defns addlSchema)
+              return . M.fromList $ zip additionalKeys (repeat . schemaGenWithFormats formatGen defns $ dereference defns addlSchema)
             _                                      -> return []
           x <- sequence $ gens <> additionalGens
           return . Object $ M.toHashMap x

test/Data/OpenApi/Schema/GeneratorSpec.hs

@@ -91,6 +91,19 @@ spec = do
     prop "MissingProperty" $ shouldNotValidate (Proxy :: Proxy MissingProperty)
     prop "WrongPropType" $ shouldNotValidate (Proxy :: Proxy WrongPropType)
 
+  describe "schemaGenWithFormats" $ do
+    it "supports custom string format generators" $ do
+      let sch = mempty
+            & type_  ?~ OpenApiString
+            & format ?~ "custom"
+
+      let formatGen fmt = case fmt of
+            "custom" -> Just (pure "custom")
+            _        -> Nothing
+
+      value <- generate $ schemaGenWithFormats formatGen mempty sch
+      value `shouldBe` String "custom"
+
 -- =============================
 -- Data types and bunk instances
 -- =============================