{-# LANGUAGE FlexibleContexts, MultiParamTypeClasses, OverloadedStrings,
ScopedTypeVariables, LambdaCase #-}
{- |
Module : Streamly.Csv
Description : Cassava support for the streamly library
Copyright : (c) Richard Warfield
License : BSD 3-clause
Maintainer : richard@litx.io
Stream CSV data in\/out using
[Cassava](http://hackage.haskell.org/package/cassava). Adapted from
[streaming-cassava](http://hackage.haskell.org/package/streaming-cassava).
For efficiency, operates on streams of strict ByteString chunks
@(i.e. IsStream t => t m ByteString)@ rather than directly on streams of Word8.
The 'chunkStream' function is useful for generating an input stream from a
'Handle'.
Example usage:
> import Streamly
> import qualified Streamly.Prelude as S
> import Streamly.Csv (decode, encode, chunkStream)
> import System.IO
> import qualified Data.Csv as Csv
> import qualified Data.ByteString as BS
> import Data.Vector (Vector)
>
> do
> h <- openFile "testfile.csv" ReadMode
> let chunks = chunkStream h (64*1024)
> recs = decode Csv.HasHeader chunks :: SerialT IO (Vector BS.ByteString)
> withFile "dest.csv" WriteMode $ \ho ->
> S.mapM_ (BS.hPut ho) $ encode Nothing recs
-}
module Streamly.Csv
( -- * Decoding
decode
, decodeWith
, decodeWithErrors
, CsvParseException (..)
, chunkStream
-- ** Named decoding
, decodeByName
, decodeByNameWith
, decodeByNameWithErrors
-- * Encoding
, encode
, encodeDefault
, encodeWith
-- ** Named encoding
, encodeByName
, encodeByNameDefault
, encodeByNameWith
-- * Re-exports
, FromRecord (..)
, FromNamedRecord (..)
, ToRecord (..)
, ToNamedRecord (..)
, DefaultOrdered (..)
, HasHeader (..)
, Header
, header
, Name
, DecodeOptions(..)
, defaultDecodeOptions
, EncodeOptions(..)
, defaultEncodeOptions
) where
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as BSL
import Streamly
import qualified Streamly.Prelude as S
import Data.Csv (DecodeOptions(..), DefaultOrdered(..),
EncodeOptions(..), FromNamedRecord(..),
FromRecord(..), Header, Name,
ToNamedRecord(..), ToRecord(..),
defaultDecodeOptions,
defaultEncodeOptions, encIncludeHeader,
header)
import Data.Csv.Incremental (HasHeader(..), HeaderParser(..),
Parser(..))
import qualified Data.Csv.Incremental as CI
import System.IO (Handle)
import Control.Exception (Exception(..))
import Control.Monad.Catch (MonadThrow(..))
import Control.Monad.IO.Class (MonadIO, liftIO)
import Data.Word (Word8)
import Data.Bifunctor (first)
import Data.Maybe (fromMaybe)
import Data.String (IsString(..))
import Data.Typeable (Typeable)
--------------------------------------------------------------------------------
-- | Use 'defaultOptions' for decoding the provided CSV.
decode :: (IsStream t, MonadAsync m, FromRecord a)
=> HasHeader
-> t m BS.ByteString
-> t m a
decode = decodeWith defaultDecodeOptions
-- | Return back a stream of values from the provided CSV, stopping at
-- the first error.
--
-- If you wish to instead ignore errors, consider using
-- 'decodeWithErrors' with 'S.mapMaybe'
--
-- Any remaining input is discarded.
decodeWith :: (IsStream t, MonadAsync m, FromRecord a)
=> DecodeOptions -> HasHeader
-> t m BS.ByteString
-> t m a
decodeWith opts hdr chunks = getValues (decodeWithErrors opts hdr chunks)
-- >>= either (throwError . fst) return
-- | Return back a stream with an attempt at type conversion, and
-- either the previous result or any overall parsing errors with the
-- remainder of the input.
decodeWithErrors :: (IsStream t, Monad m, FromRecord a, MonadThrow m)
=> DecodeOptions -> HasHeader
-> t m BS.ByteString
-> t m (Either CsvParseException a)
decodeWithErrors opts = runParser . CI.decodeWith opts
runParser :: forall t a m. (IsStream t, Monad m, MonadThrow m)
=> Parser a -> t m BS.ByteString -> t m (Either CsvParseException a)
runParser p chunked = S.concatMap fst $ S.scanlM' continue (S.nil, const p) $
S.cons BS.empty chunked
where
continue :: (t m (Either CsvParseException a), BS.ByteString -> Parser a)
-> BS.ByteString
-> m (t m (Either CsvParseException a), BS.ByteString -> Parser a)
continue (_, p) chunk =
case p chunk of
Fail bs err -> throwM (CsvParseException err)
Many es get -> return (withEach es, get)
Done es -> return (withEach es, p)
withEach = S.fromList . map (first CsvParseException)
chunkStream :: (IsStream t, MonadAsync m) => Handle -> Int -> t m BS.ByteString
chunkStream h chunkSize = loop
where
loop = S.takeWhile (not . BS.null) $
liftIO (BS.hGetSome h chunkSize) `S.consM` loop
--------------------------------------------------------------------------------
-- | Use 'defaultOptions' for decoding the provided CSV.
decodeByName :: (MonadAsync m, FromNamedRecord a)
=> SerialT m BS.ByteString -> SerialT m a
decodeByName = decodeByNameWith defaultDecodeOptions
-- | Return back a stream of values from the provided CSV, stopping at
-- the first error.
--
-- A header is required to determine the order of columns, but then
-- discarded.
--
-- If you wish to instead ignore errors, consider using
-- 'decodeByNameWithErrors' with 'S.mapMaybe'
--
-- Any remaining input is discarded.
decodeByNameWith :: (MonadAsync m, FromNamedRecord a)
=> DecodeOptions
-> SerialT m BS.ByteString -> SerialT m a
decodeByNameWith opts bs = getValues (decodeByNameWithErrors opts bs)
-- >>= either (throwError . fst) return
-- | Return back a stream with an attempt at type conversion, but
-- where the order of columns doesn't have to match the order of
-- fields of your actual type.
--
-- This requires\/assumes a header in the CSV stream, which is
-- discarded after parsing.
--
decodeByNameWithErrors :: forall m a. (Monad m, MonadThrow m, FromNamedRecord a)
=> DecodeOptions
-> SerialT m BS.ByteString
-> SerialT m (Either CsvParseException a)
decodeByNameWithErrors opts chunked = do
(p, rest) <- S.yieldM $ extractParser (const $ CI.decodeByNameWith opts) $ S.cons BS.empty chunked
runParser p rest
where
extractParser :: (BS.ByteString -> HeaderParser (Parser a))
-> SerialT m BS.ByteString
-> m (Parser a, SerialT m BS.ByteString)
extractParser p chunks = S.uncons chunks >>= \case
Just (hed, rest) ->
case p hed of
FailH bs err -> throwM (CsvParseException err)
PartialH get -> extractParser get rest
DoneH _ p -> return (p, rest)
Nothing -> throwM $ CsvParseException "Unexpected end of input stream"
-- --------------------------------------------------------------------------------
--
-- -- | Encode a stream of values with the default options.
-- --
-- -- Optionally prefix the stream with headers (the 'header' function
-- -- may be useful).
encode :: (IsStream t, ToRecord a, Monad m) => Maybe Header
-> t m a -> t m BS.ByteString
encode = encodeWith defaultEncodeOptions
-- | Encode a stream of values with the default options and a derived
-- header prefixed.
encodeDefault :: forall a t m. (IsStream t, ToRecord a, DefaultOrdered a, Monad m)
=> t m a -> t m BS.ByteString
encodeDefault = encode (Just (headerOrder (undefined :: a)))
-- | Encode a stream of values with the provided options.
--
-- Optionally prefix the stream with headers (the 'header' function
-- may be useful).
encodeWith :: (IsStream t, ToRecord a, Monad m)
=> EncodeOptions
-> Maybe Header
-> t m a
-> t m BS.ByteString
encodeWith opts mhdr = S.concatMap S.fromList
. addHeaders
. S.map enc
where
addHeaders = maybe id (S.cons . enc) mhdr
enc :: (ToRecord v) => v -> [BS.ByteString]
enc = BSL.toChunks . CI.encodeWith opts . CI.encodeRecord
--------------------------------------------------------------------------------
-- | Use the default ordering to encode all fields\/columns.
encodeByNameDefault :: forall a t m. (IsStream t, DefaultOrdered a, ToNamedRecord a, Monad m)
=> t m a -> t m BS.ByteString
encodeByNameDefault = encodeByName (headerOrder (undefined :: a))
-- | Select the columns that you wish to encode from your data
-- structure using default options (which currently includes
-- printing the header).
encodeByName :: (IsStream t, ToNamedRecord a, Monad m) => Header
-> t m a -> t m BS.ByteString
encodeByName = encodeByNameWith defaultEncodeOptions
-- | Select the columns that you wish to encode from your data
-- structure.
--
-- Header printing respects 'encIncludeheader'.
encodeByNameWith :: (IsStream t, ToNamedRecord a, Monad m) => EncodeOptions -> Header
-> t m a -> t m BS.ByteString
encodeByNameWith opts hdr = S.concatMap S.fromList
. addHeaders
. S.map enc
where
opts' = opts { encIncludeHeader = False }
addHeaders
| encIncludeHeader opts = S.cons . BSL.toChunks
. CI.encodeWith opts' . CI.encodeRecord $ hdr
| otherwise = id
enc = BSL.toChunks . CI.encodeByNameWith opts' hdr . CI.encodeNamedRecord
--------------------------------------------------------------------------------
getValues :: (IsStream t, MonadAsync m, Exception e)
=> t m (Either e a) -> t m a
getValues = S.mapM (either throwM return)
newtype CsvParseException = CsvParseException String
deriving (Eq, Show, Typeable)
instance IsString CsvParseException where
fromString = CsvParseException
instance Exception CsvParseException where
displayException (CsvParseException e) = "Error parsing csv: " ++ e