{-

This file is part of the Haskell package cassava-streams. It is
subject to the license terms in the LICENSE file found in the
top-level directory of this distribution and at
git://pmade.com/cassava-streams/LICENSE. No part of cassava-streams
package, including this file, may be copied, modified, propagated, or
distributed except according to the terms contained in the LICENSE
file.

-}

--------------------------------------------------------------------------------
module System.IO.Streams.Csv.Encode
       ( encodeStream
       , encodeStreamWith
       , encodeStreamByName
       , encodeStreamByNameWith
       ) where

--------------------------------------------------------------------------------
import Control.Monad (when)
import Data.ByteString (ByteString)
import qualified Data.ByteString.Lazy as BL
import Data.Csv
import Data.IORef
import System.IO.Streams (OutputStream, makeOutputStream)
import qualified System.IO.Streams as Streams

--------------------------------------------------------------------------------
-- | Create a new @OutputStream@ that can be fed @ToRecord@ values
-- which are converted to CSV.  The records are encoded into
-- @ByteString@s and passed on to the given downstream @OutputStream@.
--
-- Equivalent to @encodeStreamWith defaultEncodeOptions@.
encodeStream :: ToRecord a
             => OutputStream ByteString -- ^ Downstream.
             -> IO (OutputStream a)     -- ^ New @OutputStream@.
encodeStream :: OutputStream ByteString -> IO (OutputStream a)
encodeStream = EncodeOptions -> OutputStream ByteString -> IO (OutputStream a)
forall a.
ToRecord a =>
EncodeOptions -> OutputStream ByteString -> IO (OutputStream a)
encodeStreamWith EncodeOptions
defaultEncodeOptions

--------------------------------------------------------------------------------
-- | Create a new @OutputStream@ that can be fed @ToRecord@ values
-- which are converted to CSV.  The records are encoded into
-- @ByteString@s and passed on to the given downstream @OutputStream@.
encodeStreamWith :: ToRecord a
                 => EncodeOptions           -- ^ Encoding options.
                 -> OutputStream ByteString -- ^ Downstream.
                 -> IO (OutputStream a)     -- ^ New @OutputStream@.
encodeStreamWith :: EncodeOptions -> OutputStream ByteString -> IO (OutputStream a)
encodeStreamWith EncodeOptions
opts OutputStream ByteString
output = do
  IORef EncodeOptions
ref <- EncodeOptions -> IO (IORef EncodeOptions)
forall a. a -> IO (IORef a)
newIORef EncodeOptions
opts
  (Maybe a -> IO ()) -> IO (OutputStream a)
forall a. (Maybe a -> IO ()) -> IO (OutputStream a)
makeOutputStream ((EncodeOptions -> [a] -> ByteString)
-> IORef EncodeOptions
-> OutputStream ByteString
-> Maybe a
-> IO ()
forall a.
(EncodeOptions -> [a] -> ByteString)
-> IORef EncodeOptions
-> OutputStream ByteString
-> Maybe a
-> IO ()
dispatch EncodeOptions -> [a] -> ByteString
forall a. ToRecord a => EncodeOptions -> [a] -> ByteString
encodeWith IORef EncodeOptions
ref OutputStream ByteString
output)

--------------------------------------------------------------------------------
-- | Create a new @OutputStream@ which can be fed @ToNamedRecord@
-- values that will be converted into CSV.  The records are encoded
-- into @ByteString@s and passed on to the given downstream
-- @OutputStream@.
--
-- Equivalent to @encodeStreamByNameWith defaultEncodeOptions@.
encodeStreamByName :: ToNamedRecord a
                   => Header                   -- ^ CSV Header.
                   -> OutputStream ByteString  -- ^ Downstream.
                   -> IO (OutputStream a)      -- ^ New @OutputStream@.
encodeStreamByName :: Header -> OutputStream ByteString -> IO (OutputStream a)
encodeStreamByName = EncodeOptions
-> Header -> OutputStream ByteString -> IO (OutputStream a)
forall a.
ToNamedRecord a =>
EncodeOptions
-> Header -> OutputStream ByteString -> IO (OutputStream a)
encodeStreamByNameWith EncodeOptions
defaultEncodeOptions

--------------------------------------------------------------------------------
-- | Create a new @OutputStream@ which can be fed @ToNamedRecord@
-- values that will be converted into CSV.  The records are encoded
-- into @ByteString@s and passed on to the given downstream
-- @OutputStream@.
encodeStreamByNameWith :: ToNamedRecord a
                       => EncodeOptions            -- ^ Encoding options.
                       -> Header                   -- ^ CSV Header.
                       -> OutputStream ByteString  -- ^ Downstream.
                       -> IO (OutputStream a)      -- ^ New @OutputStream@.
encodeStreamByNameWith :: EncodeOptions
-> Header -> OutputStream ByteString -> IO (OutputStream a)
encodeStreamByNameWith EncodeOptions
opts Header
hdr OutputStream ByteString
output = do
  IORef EncodeOptions
ref <- EncodeOptions -> IO (IORef EncodeOptions)
forall a. a -> IO (IORef a)
newIORef EncodeOptions
opts
  (Maybe a -> IO ()) -> IO (OutputStream a)
forall a. (Maybe a -> IO ()) -> IO (OutputStream a)
makeOutputStream ((Maybe a -> IO ()) -> IO (OutputStream a))
-> (Maybe a -> IO ()) -> IO (OutputStream a)
forall a b. (a -> b) -> a -> b
$ (EncodeOptions -> [a] -> ByteString)
-> IORef EncodeOptions
-> OutputStream ByteString
-> Maybe a
-> IO ()
forall a.
(EncodeOptions -> [a] -> ByteString)
-> IORef EncodeOptions
-> OutputStream ByteString
-> Maybe a
-> IO ()
dispatch (EncodeOptions -> Header -> [a] -> ByteString
forall a.
ToNamedRecord a =>
EncodeOptions -> Header -> [a] -> ByteString
`encodeByNameWith` Header
hdr) IORef EncodeOptions
ref OutputStream ByteString
output

--------------------------------------------------------------------------------
-- | Encode records, ensuring that the header is written no more than once.
dispatch :: (EncodeOptions -> [a] -> BL.ByteString) -- ^ Encoding function.
         -> IORef EncodeOptions                     -- ^ Encoding options.
         -> OutputStream ByteString                 -- ^ Downstream.
         -> Maybe a                                 -- ^ Record to write.
         -> IO ()
dispatch :: (EncodeOptions -> [a] -> ByteString)
-> IORef EncodeOptions
-> OutputStream ByteString
-> Maybe a
-> IO ()
dispatch EncodeOptions -> [a] -> ByteString
_   IORef EncodeOptions
_   OutputStream ByteString
output Maybe a
Nothing  = Maybe ByteString -> OutputStream ByteString -> IO ()
forall a. Maybe a -> OutputStream a -> IO ()
Streams.write Maybe ByteString
forall a. Maybe a
Nothing OutputStream ByteString
output
dispatch EncodeOptions -> [a] -> ByteString
enc IORef EncodeOptions
ref OutputStream ByteString
output (Just a
x) = do
  EncodeOptions
opts <- IORef EncodeOptions -> IO EncodeOptions
forall a. IORef a -> IO a
readIORef IORef EncodeOptions
ref
  Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (EncodeOptions -> Bool
encIncludeHeader EncodeOptions
opts) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ IORef EncodeOptions -> EncodeOptions -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef EncodeOptions
ref (EncodeOptions
opts {encIncludeHeader :: Bool
encIncludeHeader = Bool
False})
  ByteString -> OutputStream ByteString -> IO ()
Streams.writeLazyByteString (EncodeOptions -> [a] -> ByteString
enc EncodeOptions
opts [a
x]) OutputStream ByteString
output