nested.hs

{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE NoImplicitPrelude #-}

import BasePrelude
import Data.Text (Text)
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import qualified Data.HashMap.Strict as HM
import qualified Data.Vector as V
import Data.Csv

-- Here are our nested datatypes. We want CSV to look like this:
--
--     name,age,address_country,address_city
--     Alice,47,Germany,Berlin
--     Bob,32,Austria,Graz

data Person = Person
  { name    :: Text
  , age     :: Int
  , address :: Address
  } deriving (Eq, Show, Generic)

data Address = Address
  { country :: Text
  , city    :: Text
  } deriving (Eq, Show, Generic)

testData :: [Person]
testData =
  [ Person "Alice" 47 (Address "Germany" "Berlin")
  , Person "Bob" 32 (Address "Austria" "Graz")
  ]

-- The instances for 'Address' can just be autoderived. The field
-- names will be "country" and "city" instead of "address_country"
-- and "address_city", but actually it's exactly what we want.
-- The instance for 'Person' will be responsible for adding the
-- prefix.

instance ToNamedRecord Address
instance FromNamedRecord Address
instance DefaultOrdered Address

-- For the 'ToNamedRecord Person' instance, we'll construct the
-- hashmap by hand, distinguishing between fields that should
-- be encoded normally and fields that should be expanded into
-- several columns.

-- | Add a prefix to all field names.
prefix :: ByteString -> NamedRecord -> NamedRecord
prefix p r = HM.fromList [(p <> k, v) | (k, v) <- HM.toList r]

instance ToNamedRecord Person where
  toNamedRecord Person{..} =
    -- Normal fields
    namedRecord
      [ "name" .= name
      , "age"  .= age
      ] <>
    -- 'address'
    prefix "address_" (toNamedRecord address)
    -- more fields could be here (don't forget the underscores!
    -- it might be safer to add underscore handling directly
    -- to 'prefix')

-- The 'FromNamedRecord Person' instance is slightly more
-- involved; the hashmap has to be traversed once for each
-- nested field. If speed is important, splitting the hashmap
-- into buckets based on prefix would be faster. If speed is
-- *really* important, use 'FromRecord' and 'ToRecord'.

-- | Find all fields with given prefix and strip it out.
select :: ByteString -> NamedRecord -> NamedRecord
select p r = HM.fromList [(k, v) | (pk, v) <- HM.toList r
                                 , Just k <- [BS.stripPrefix p pk]]

instance FromNamedRecord Person where
  parseNamedRecord r = do
    name    <- r .: "name"
    age     <- r .: "age"
    address <- parseNamedRecord (select "address_" r)
    pure Person{..}

-- The 'DefaultOrdered Person' instance is very similar to the
-- 'ToNamedRecord Person' instance.

instance DefaultOrdered Person where
  headerOrder _ =
    -- Normal fields
    V.fromList ["name", "age"] <>
    -- 'address'
    fmap ("address_" <>) (headerOrder @Address undefined)

-- Test:
--
--     > encodeDefaultOrderedByName testData
--     "name,age,address_country,address_city\r\n\
--     \Alice,47,Germany,Berlin\r\n\
--     \Bob,32,Austria,Graz\r\n"
--
--     > let Right (_, parsed) = decodeByName it
--     > parsed == V.fromList testData
--     True