dataframe-2.3.0.0: tests/GenDataFrame.hs

{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeApplications #-}
{-# OPTIONS_GHC -Wno-orphans #-}

module GenDataFrame where

import qualified Data.Map as M
import qualified Data.Text as T
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as VU
import DataFrame.Internal.Column
import DataFrame.Internal.DataFrame
import Test.QuickCheck

genColumn :: Int -> Gen Column
genColumn len =
    oneof
        [ BoxedColumn Nothing . V.fromList <$> vectorOf len (arbitrary @Int)
        , UnboxedColumn Nothing . VU.fromList <$> vectorOf len (arbitrary @Double)
        , fromVector . V.fromList <$> vectorOf len (arbitrary @(Maybe Int))
        ]

genDataFrame :: Gen DataFrame
genDataFrame = do
    numRows <- choose (100, 1000)
    numCols <- choose (1, 10)
    colNames <- V.fromList <$> vectorOf numCols genUniqueColName
    cols <- V.fromList <$> vectorOf numCols (genColumn numRows)
    let indices = M.fromList $ zip (V.toList colNames) [0 ..]
    pure $
        DataFrame
            { columns = cols
            , columnIndices = indices
            , dataframeDimensions = (numRows, numCols)
            , derivingExpressions = M.empty
            }

genUniqueColName :: Gen T.Text
genUniqueColName = T.pack <$> listOf1 (elements ['a' .. 'z'])

instance Arbitrary DataFrame where
    arbitrary = genDataFrame
    shrink _df = []