Conduit.hs

{-# LANGUAGE BangPatterns          #-}
{-# LANGUAGE CPP                   #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE UndecidableInstances  #-}

module Data.CSV.Conduit
    (

    -- * Main Interface
      decodeCSV
    , readCSVFile
    , writeCSVFile
    , transformCSV
    , transformCSV'
    , mapCSVFile
    , writeHeaders
    , writeHeadersOrdered

    -- Types
    , CSV (..)
    , CSVSettings (..)
    , defCSVSettings
    , MapRow
    , OrderedMapRow
    , Row

    -- * Re-exported For Convenience
    , runResourceT
    ) where

-------------------------------------------------------------------------------
import           Control.Exception
import           Control.Monad.Catch.Pure           (runCatchT, CatchT)
import           Control.Monad.IO.Class             (MonadIO (liftIO))
import           Control.Monad.Primitive
import           Control.Monad.ST
import           Control.Monad.Trans.Class          (MonadTrans(lift))
import           Control.Monad.Trans.Except         (ExceptT(..), runExceptT)
import           Control.Monad.Trans.Resource       (MonadResource, MonadThrow,
                                                     runResourceT)
import           Data.Attoparsec.Types              (Parser)
import qualified Data.ByteString                    as B
import           Data.ByteString.Char8              (ByteString)
import qualified Data.ByteString.Char8              as B8
import           Data.ByteString.Internal           (c2w)
import           Data.Conduit
import           Data.Conduit.Attoparsec
import           Data.Conduit.Binary                (sinkFile, sinkIOHandle,
                                                     sourceFile)
import qualified Data.Conduit.List                  as C
import qualified Data.Map                           as M
import qualified Data.Map.Ordered                   as MO
import           Data.String
import           Data.Text                          (Text)
import qualified Data.Text                          as T
import qualified Data.Text.Encoding                 as T
import qualified Data.Vector                        as V
import qualified Data.Vector.Generic                as GV
import qualified Data.Vector.Generic.Mutable        as GMV
import           Data.Void                           as Void
import           System.IO
-------------------------------------------------------------------------------
import           Data.CSV.Conduit.Conversion        (FromNamedRecord (..),
                                                     FromNamedRecordOrdered (..),
                                                     Named (..),
                                                     NamedOrdered (..),
                                                     ToNamedRecord (..),
                                                     ToNamedRecordOrdered (..),
                                                     runParser)
import qualified Data.CSV.Conduit.Parser.ByteString as BSP
import qualified Data.CSV.Conduit.Parser.Text       as TP
import           Data.CSV.Conduit.Types
-------------------------------------------------------------------------------


-------------------------------------------------------------------------------
-- | Represents types 'r' that are CSV-like and can be converted
-- to/from an underlying stream of type 's'. There is nothing scary
-- about the type:
--
-- @s@ represents stream types that can be converted to\/from CSV rows.
-- Examples are 'ByteString', 'Text' and 'String'.
--
-- @r@ represents the target CSV row representations that this library
-- can work with. Examples are the 'Row' types, the 'Record' type and
-- the 'MapRow' family of types. We can also convert directly to
-- complex Haskell types using the 'Data.CSV.Conduit.Conversion'
-- module that was borrowed from the cassava package, which was itself
-- inspired by the aeson package.
--
--
-- Example #1: Basics Using Convenience API
--
-- >import Data.Conduit
-- >import Data.Conduit.Binary
-- >import Data.Conduit.List as CL
-- >import Data.CSV.Conduit
-- >
-- >myProcessor :: Conduit (Row Text) m (Row Text)
-- >myProcessor = CL.map reverse
-- >
-- >test = runResourceT $
-- >  transformCSV defCSVSettings
-- >               (sourceFile "input.csv")
-- >               myProcessor
-- >               (sinkFile "output.csv")
--
--
-- Example #2: Basics Using Conduit API
--
-- >import Data.Conduit
-- >import Data.Conduit.Binary
-- >import Data.CSV.Conduit
-- >
-- >myProcessor :: Conduit (MapRow Text) m (MapRow Text)
-- >myProcessor = undefined
-- >
-- >test = runResourceT $ runConduit $
-- >  sourceFile "test/BigFile.csv" .|
-- >  intoCSV defCSVSettings .|
-- >  myProcessor .|
-- >  (writeHeaders defCSVSettings >> fromCSV defCSVSettings) .|
-- >  sinkFile "test/BigFileOut.csv"
class CSV s r where

  -----------------------------------------------------------------------------
  -- | Convert a CSV row into strict ByteString equivalent.
  rowToStr :: CSVSettings -> r -> s

  -----------------------------------------------------------------------------
  -- | Turn a stream of 's' into a stream of CSV row type. An example
  -- would be parsing a ByteString stream as rows of 'MapRow' 'Text'.
  intoCSV :: (MonadThrow m) => CSVSettings -> ConduitT s r m ()

  -----------------------------------------------------------------------------
  -- | Turn a stream of CSV row type back into a stream of 's'. An
  -- example would be rendering a stream of 'Row' 'ByteString' rows as
  -- 'Text'.
  fromCSV :: Monad m => CSVSettings -> ConduitT r s m ()





------------------------------------------------------------------------------
-- | 'Row' instance using 'ByteString'
instance CSV ByteString (Row ByteString) where
  rowToStr s !r =
    let
      sep = B.pack [c2w (csvSep s)]
      wrapField !f = case csvQuoteCharAndStyle s of
        Just (x, quoteEmpty) ->
          case quoteEmpty == DoQuoteEmpty || B8.length f /= 0 of
            True -> (x `B8.cons` escape x f) `B8.snoc` x
            False -> f
        Nothing   -> f
      escape c str = B8.intercalate (B8.pack [c,c]) $ B8.split c str
    in B.intercalate sep . map wrapField $ r

  intoCSV set = intoCSVRow (BSP.row set)
  fromCSV set = fromCSVRow set


------------------------------------------------------------------------------
-- | 'Row' instance using 'Text'
instance CSV Text (Row Text) where
  rowToStr s !r =
    let
      sep = T.pack [csvSep s]
      wrapField !f = case csvQuoteCharAndStyle s of
        Just (x, quoteEmpty) -> case quoteEmpty == DoQuoteEmpty || not (T.null f) of
          True -> x `T.cons` escape x f `T.snoc` x
          False -> f
        Nothing -> f
      escape c str = T.intercalate (T.pack [c,c]) $ T.split (== c) str
    in T.intercalate sep . map wrapField $ r

  intoCSV set = intoCSVRow (TP.row set)
  fromCSV set = fromCSVRow set


-------------------------------------------------------------------------------
-- | 'Row' instance using 'Text' based on 'ByteString' stream
instance CSV ByteString (Row Text) where
    rowToStr s r = T.encodeUtf8 $ rowToStr s r
    intoCSV set = intoCSV set .| C.map (map T.decodeUtf8)
    fromCSV set = fromCSV set .| C.map T.encodeUtf8



-------------------------------------------------------------------------------
-- | 'Row' instance using 'String' based on 'ByteString' stream.
-- Please note this uses the ByteString operations underneath and has
-- lots of unnecessary overhead. Included for convenience.
instance CSV ByteString (Row String) where
    rowToStr s r = rowToStr s $ map B8.pack r
    intoCSV set = intoCSV set .| C.map (map B8.unpack)
    fromCSV set = C.map (map B8.pack) .| fromCSV set


-- | Support for parsing rows in the 'Vector' form.
instance (CSV s (Row s)) => CSV s (V.Vector s) where
    rowToStr s r = rowToStr s . V.toList $ r
    intoCSV set = intoCSV set .| C.map (V.fromList)
    fromCSV set = C.map (V.toList) .| fromCSV set



-------------------------------------------------------------------------------
fromCSVRow :: (Monad m, IsString s, CSV s r)
           => CSVSettings -> ConduitT r s m ()
fromCSVRow set = awaitForever $ \row -> mapM_ yield [rowToStr set row, "\n"]



-------------------------------------------------------------------------------
intoCSVRow :: (MonadThrow m, AttoparsecInput i) => Parser i (Maybe o) -> ConduitT i o m ()
intoCSVRow p = parse .| puller
  where
    parse = {-# SCC "conduitParser_p" #-} conduitParser p
    puller = {-# SCC "puller" #-}
      awaitForever $ \ (_, mrow) -> maybe (return ()) yield mrow


-------------------------------------------------------------------------------
-- | Generic 'MapRow' instance; any stream type with a 'Row' instance
-- automatically gets a 'MapRow' instance.
instance (CSV s (Row s'), Ord s', IsString s) => CSV s (MapRow s') where
  rowToStr s r = rowToStr s . M.elems $ r
  intoCSV set = intoCSVMap set
  fromCSV set = fromCSVMap set

instance (CSV s (Row s'), Ord s', IsString s) => CSV s (OrderedMapRow s') where
  rowToStr s r = rowToStr s . (map snd . MO.assocs) $ r
  intoCSV set = intoCSVMapOrdered set
  fromCSV set = fromCSVMapOrdered set


-------------------------------------------------------------------------------
intoCSVMap :: (Ord a, MonadThrow m, CSV s [a])
           => CSVSettings -> ConduitT s (MapRow a) m ()
intoCSVMap set = intoCSV set .| (headers >>= converter)
  where
    headers = do
      mrow <- await
      case mrow of
        Nothing -> return []
        Just [] -> headers
        Just hs -> return hs
    converter hs = awaitForever $ yield . toMapCSV hs
    toMapCSV !hs !fs = M.fromList $ zip hs fs

intoCSVMapOrdered :: (Ord a, MonadThrow m, CSV s [a])
           => CSVSettings -> ConduitM s (OrderedMapRow a) m ()
intoCSVMapOrdered set = intoCSV set .| (headers >>= converter)
  where
    headers = do
      mrow <- await
      case mrow of
        Nothing -> return []
        Just [] -> headers
        Just hs -> return hs
    converter hs = awaitForever $ yield . toMapCSV hs
    toMapCSV !hs !fs = MO.fromList $ zip hs fs


-- | Conversion of stream directly to/from a custom complex haskell
-- type.
instance (FromNamedRecord a, ToNamedRecord a, CSV s (MapRow ByteString)) =>
    CSV s (Named a) where
    rowToStr s a = rowToStr s . toNamedRecord . getNamed $ a
    intoCSV set = intoCSV set .| C.mapMaybe go
        where
          go x = either (const Nothing) (Just . Named) $
                 runParser (parseNamedRecord x)

    fromCSV set = C.map go .| fromCSV set
        where
          go = toNamedRecord . getNamed

instance (FromNamedRecordOrdered a, ToNamedRecordOrdered a, CSV s (OrderedMapRow ByteString)) =>
    CSV s (NamedOrdered a) where
    rowToStr s a = rowToStr s . toNamedRecordOrdered . getNamedOrdered $ a
    intoCSV set = intoCSV set .| C.mapMaybe go
        where
          go x = either (const Nothing) (Just . NamedOrdered) $
                 runParser (parseNamedRecordOrdered x)

    fromCSV set = C.map go .| fromCSV set
        where
          go = toNamedRecordOrdered . getNamedOrdered


-------------------------------------------------------------------------------
fromCSVMap :: (Monad m, IsString s, CSV s [a])
           => CSVSettings -> ConduitT (M.Map k a) s m ()
fromCSVMap set = awaitForever push
  where
    push r = mapM_ yield [rowToStr set (M.elems r), "\n"]

fromCSVMapOrdered :: (Monad m, IsString s, CSV s [a])
                  => CSVSettings -> ConduitM (MO.OMap k a) s m ()
fromCSVMapOrdered set = awaitForever push
  where
    push r = mapM_ yield [rowToStr set (map snd $ MO.assocs r), "\n"]


-------------------------------------------------------------------------------
-- | Write headers AND the row into the output stream, once. If you
-- don't call this while using 'MapRow' family of row types, then your
-- resulting output will NOT have any headers in it.
--
-- Usage: Just chain this using the 'Monad' instance in your pipeline:
--
-- > runConduit $ ... .| writeHeaders settings >> fromCSV settings .| sinkFile "..."
writeHeaders
    :: (Monad m, CSV s (Row r), IsString s)
    => CSVSettings
    -> ConduitT (MapRow r) s m ()
writeHeaders set = do
  mrow <- await
  case mrow of
    Nothing -> return ()
    Just row -> mapM_ yield [ rowToStr set (M.keys row)
                            , "\n"
                            , rowToStr set (M.elems row)
                            , "\n" ]

writeHeadersOrdered
    :: (Monad m, CSV s (Row r), IsString s)
    => CSVSettings
    -> ConduitM (OrderedMapRow r) s m ()
writeHeadersOrdered set = do
  mrow <- await
  case mrow of
    Nothing -> return ()
    Just row -> mapM_ yield [ rowToStr set (map fst $ MO.assocs row)
                            , "\n"
                            , rowToStr set (map snd $ MO.assocs row)
                            , "\n" ]


                          ---------------------------
                          -- Convenience Functions --
                          ---------------------------


-------------------------------------------------------------------------------
-- | Read the entire contents of a CSV file into memory.
readCSVFile
    :: (MonadIO m, CSV ByteString a)
    => CSVSettings -- ^ Settings to use in deciphering stream
    -> FilePath    -- ^ Input file
    -> m (V.Vector a)
readCSVFile set fp = liftIO . runResourceT $ runConduit $ sourceFile fp .| intoCSV set .| transPipe lift (sinkVector growthFactor)
  where
    growthFactor = 10


-------------------------------------------------------------------------------
-- | A simple way to decode a CSV string. Don't be alarmed by the
-- polymorphic nature of the signature. 's' is the type for the string
-- and 'v' is a kind of 'Vector' here.
--
-- For example for 'ByteString':
--
-- >>> s <- LB.readFile "my.csv"
-- >>> decodeCSV defCSVSettings s :: Either SomeException (Vector (Vector ByteString))
--
-- will work as long as the data is comma separated.
decodeCSV
    :: forall v a s. (GV.Vector v a, CSV s a)
    => CSVSettings
    -> s
    -> Either SomeException (v a)
decodeCSV set bs = runST $ runExceptT pipeline
  where
    src :: ConduitT () s (ExceptT SomeException (ST s1)) ()
    src = C.sourceList [bs]
    csvConvert :: ConduitT s a (ExceptT SomeException (ST s1)) ()
    csvConvert = transPipe (ExceptT . runCatchT) csvConvert'
    csvConvert' :: ConduitT s a (CatchT (ST s1)) ()
    csvConvert' = intoCSV set
    growthFactor = 10
    sink :: ConduitT a Void.Void (ExceptT SomeException (ST s1)) (v a)
    sink = sinkVector growthFactor
    pipeline :: ExceptT SomeException (ST s1) (v a)
    pipeline = runConduit (src .| csvConvert .| sink)



-------------------------------------------------------------------------------
-- | Write CSV data into file. As we use a 'ByteString' sink, you'll
-- need to get your data into a 'ByteString' stream type.
writeCSVFile
  :: (CSV ByteString a)
  => CSVSettings
  -- ^ CSV Settings
  -> FilePath
  -- ^ Target file
  -> IOMode
  -- ^ Write vs. append mode
  -> [a]
  -- ^ List of rows
  -> IO ()
writeCSVFile set fo fmode rows = runResourceT $ runConduit $ do
  C.sourceList rows .| fromCSV set .|
    sinkIOHandle (openFile fo fmode)


-------------------------------------------------------------------------------
-- | Map over the rows of a CSV file. Provided for convenience for
-- historical reasons.
--
-- An easy way to run this function would be 'runResourceT' after
-- feeding it all the arguments.
mapCSVFile
    :: ( MonadResource m
       , CSV ByteString a
       , CSV ByteString b
# if MIN_VERSION_resourcet(1,2,0)
       , MonadThrow m
#endif
       )
      => CSVSettings
      -- ^ Settings to use both for both input and output
      -> (a -> [b])
      -- ^ A mapping function
      -> FilePath
      -- ^ Input file
      -> FilePath
      -- ^ Output file
      -> m ()
mapCSVFile set f fi fo =
  transformCSV set (sourceFile fi) (C.concatMap f) (sinkFile fo)


-------------------------------------------------------------------------------
-- | Like transformCSV' but uses the same settings for both input and
-- output.
transformCSV
    :: (MonadThrow m, CSV s a, CSV s' b)
    => CSVSettings
    -- ^ Settings to be used for both input and output
    -> ConduitT () s m ()
    -- ^ A raw stream data source. Ex: 'sourceFile inFile'
    -> ConduitT a b m ()
    -- ^ A transforming conduit
    -> ConduitT s' Void.Void m ()
    -- ^ A raw stream data sink. Ex: 'sinkFile outFile'
    -> m ()
transformCSV set = transformCSV' set set


-------------------------------------------------------------------------------
-- | General purpose CSV transformer. Apply a list-like processing
-- function from 'Data.Conduit.List' to the rows of a CSV stream. You
-- need to provide a stream data source, a transformer and a stream
-- data sink.
--
-- An easy way to run this function would be 'runResourceT' after
-- feeding it all the arguments.
--
-- Example - map a function over the rows of a CSV file:
--
-- > transformCSV setIn setOut (sourceFile inFile) (C.map f) (sinkFile outFile)
transformCSV'
    :: (MonadThrow m, CSV s a, CSV s' b)
    => CSVSettings
    -- ^ Settings to be used for input
    -> CSVSettings
    -- ^ Settings to be used for output
    -> ConduitT () s m ()
    -- ^ A raw stream data source. Ex: 'sourceFile inFile'
    -> ConduitT a b m ()
    -- ^ A transforming conduit
    -> ConduitT s' Void.Void m ()
    -- ^ A raw stream data sink. Ex: 'sinkFile outFile'
    -> m ()
transformCSV' setIn setOut source c sink = runConduit $
    source .|
    intoCSV setIn .|
    c .|
    fromCSV setOut .|
    sink




                              ------------------
                              -- Vector Utils --
                              ------------------



-------------------------------------------------------------------------------
-- | An efficient sink that incrementally grows a vector from the input stream
sinkVector :: (PrimMonad m, GV.Vector v a) => Int -> ConduitT a o m (v a)
sinkVector by = do
    v <- lift $ GMV.new by
    go 0 v
  where
    -- i is the index of the next element to be written by go
    -- also exactly the number of elements in v so far
    go i v = do
        res <- await
        case res of
          Nothing -> do
            v' <- lift $ GV.freeze $ GMV.slice 0 i v
            return $! v'
          Just x -> do
            v' <- case GMV.length v == i of
                    True  -> lift $ GMV.grow v by
                    False -> return v
            lift $ GMV.write v' i x
            go (i+1) v'