faster-megaparsec-0.1.2.0: test/GenCsv.hs

{-# LANGUAGE OverloadedStrings #-}
module GenCsv (CSVStructure(..),genCSVStructure,genCSVFile) where
import qualified Test.QuickCheck as Q
import Data.Text (Text)
import qualified Data.Text as T
import GenLanguage (genAlphaChar)


genTextField :: Q.Gen Text
genTextField = fmap T.pack (Q.listOf genAlphaChar)

genNumField :: Q.Gen Text
genNumField = fmap (T.pack . (show :: Int -> String)) Q.arbitrary

genBoolField :: Q.Gen Text
genBoolField = fmap (T.pack . (show :: Bool -> String)) Q.arbitrary

data CSVStructure = CSVStructure {
    tableWidth :: Int,
    fooCol :: Int,  -- ^ an 'Int' column
    barCol :: Int,  -- ^ a 'Bool' column
    columnTitles :: [Text] -- we include a title row so that we can test against cassava's NamedRecord decoding
    }

csvHeader :: CSVStructure -> Text
csvHeader = T.intercalate "," . columnTitles

-- | generate a CSV file structure:
-- two columns named @foo@ and @bar@ 
-- at random positions in a table of random width.
genCSVStructure :: Q.Gen CSVStructure
genCSVStructure = do
    width <- Q.chooseInt (5,50)
    foo <- Q.chooseInt (0,width-2)
    bar <-  Q.chooseInt (foo+1,width-1)
    beforeFoo <- sequence (replicate foo genTextField)
    betweenFooBar <- sequence (replicate (bar-foo-1) genTextField)
    afterBar <- sequence (replicate (width-1-bar) genTextField)
    return $ CSVStructure {
        tableWidth = width,
        fooCol = foo,
        barCol = bar,
        columnTitles = beforeFoo <> ["foo"] <> betweenFooBar <> ["bar"] <> afterBar
        }

-- | generate a row of given structure
genCSVrow :: CSVStructure -> Q.Gen Text
genCSVrow c = (fmap (T.intercalate ",") . sequence) (
    (replicate (fooCol c) genTextField) ++ 
    [genNumField] ++ 
    (replicate (barCol c - fooCol c - 1) genTextField) ++
    [genBoolField] ++
    (replicate (tableWidth c - 1 - barCol c) genTextField) 
    )

-- | generate a csv table of given structure 
-- and number of data rows
genCSVFile :: CSVStructure -> Int -> Q.Gen Text
genCSVFile c rows = fmap 
    (T.unlines . ((csvHeader c):)) 
    (sequence (replicate rows (genCSVrow c)))