@@ -0,0 +1,166 @@
1+{-# LANGUAGE DeriveDataTypeable #-}
2+{-# LANGUAGE ScopedTypeVariables #-}
3+
4+{-
5+
6+This file is part of the Haskell package cassava-streams. It is
7+subject to the license terms in the LICENSE file found in the
8+top-level directory of this distribution and at
9+git://pmade.com/cassava-streams/LICENSE. No part of cassava-streams
10+package, including this file, may be copied, modified, propagated, or
11+distributed except according to the terms contained in the LICENSE
12+file.
13+
14+-}
15+
16+--------------------------------------------------------------------------------
17+module System.IO.Streams.Csv.Decode
18+ ( StreamDecodingError (..)
19+ , decodeStream
20+ , decodeStreamWith
21+ , decodeStreamByName
22+ , decodeStreamByNameWith
23+ , onlyValidRecords
24+ ) where
25+
26+--------------------------------------------------------------------------------
27+import Control.Exception
28+import Data.ByteString (ByteString)
29+import qualified Data.ByteString as BS
30+import Data.Csv hiding (Parser, decodeWith, decodeByNameWith)
31+import Data.Csv.Incremental
32+import Data.IORef
33+import Data.Typeable
34+import System.IO.Streams (InputStream, makeInputStream)
35+import qualified System.IO.Streams as Streams
36+
37+--------------------------------------------------------------------------------
38+-- | Exception thrown when stream decoding cannot continue due to an
39+-- error.
40+newtype StreamDecodingError = StreamDecodingError String
41+ deriving (Typeable, Show)
42+
43+instance Exception StreamDecodingError
44+
45+--------------------------------------------------------------------------------
46+-- | Create an @InputStream@ which decodes CSV records from the given
47+-- upstream data source.
48+--
49+-- Equivalent to @decodeStreamWith defaultDecodeOptions@.
50+decodeStream :: (FromRecord a)
51+ => HasHeader
52+ -- ^ Whether to skip a header or not.
53+ -> InputStream ByteString
54+ -- ^ Upstream.
55+ -> IO (InputStream (Either String a))
56+ -- ^ An @InputStream@ which produces records.
57+decodeStream = decodeStreamWith defaultDecodeOptions
58+
59+--------------------------------------------------------------------------------
60+-- | Create an @InputStream@ which decodes CSV records from the given
61+-- upstream data source.
62+decodeStreamWith :: (FromRecord a)
63+ => DecodeOptions
64+ -- ^ CSV decoding options.
65+ -> HasHeader
66+ -- ^ Whether to skip a header or not.
67+ -> InputStream ByteString
68+ -- ^ Upstream.
69+ -> IO (InputStream (Either String a))
70+ -- ^ An @InputStream@ which produces records.
71+decodeStreamWith ops hdr input = do
72+ queue <- newIORef []
73+ parser <- newIORef $ Just (decodeWith ops hdr)
74+ makeInputStream (dispatch queue parser input)
75+
76+--------------------------------------------------------------------------------
77+-- | Create an @InputStream@ which decodes CSV records from the given
78+-- upstream data source. Data should be preceded by a header.
79+--
80+-- Equivalent to @decodeStreamByNameWith defaultDecodeOptions@.
81+decodeStreamByName :: (FromNamedRecord a)
82+ => InputStream ByteString
83+ -- ^ Upstream.
84+ -> IO (InputStream (Either String a))
85+ -- ^ An @InputStream@ which produces records.
86+decodeStreamByName = decodeStreamByNameWith defaultDecodeOptions
87+
88+--------------------------------------------------------------------------------
89+-- | Create an @InputStream@ which decodes CSV records from the given
90+-- upstream data source. Data should be preceded by a header.
91+decodeStreamByNameWith :: (FromNamedRecord a)
92+ => DecodeOptions
93+ -- ^ CSV decoding options.
94+ -> InputStream ByteString
95+ -- ^ Upstream.
96+ -> IO (InputStream (Either String a))
97+ -- ^ An @InputStream@ which produces records.
98+decodeStreamByNameWith ops input = go (decodeByNameWith ops) where
99+ -- Dispatch on the HeaderParser type.
100+ go (FailH _ e) = bomb e
101+ go (PartialH f) = Streams.read input >>= go . maybe (f BS.empty) f
102+ go (DoneH _ p) = do
103+ queue <- newIORef []
104+ parser <- newIORef (Just p)
105+ makeInputStream (dispatch queue parser input)
106+
107+--------------------------------------------------------------------------------
108+-- | Creates a new @InputStream@ which only sends valid CSV records
109+-- downstream. The first invalid record will throw an exception.
110+onlyValidRecords :: InputStream (Either String a)
111+ -- ^ Upstream.
112+ -> IO (InputStream a)
113+ -- ^ An @InputStream@ which only produces valid
114+ -- records.
115+onlyValidRecords input = makeInputStream $ do
116+ upstream <- Streams.read input
117+
118+ case upstream of
119+ Nothing -> return Nothing
120+ Just (Left err) -> bomb ("invalid CSV row: " ++ err)
121+ Just (Right x) -> return (Just x)
122+
123+--------------------------------------------------------------------------------
124+-- | Internal function which feeds data to the CSV parser.
125+dispatch :: forall a. IORef [Either String a]
126+ -- ^ List of queued CSV records.
127+ -> IORef (Maybe (Parser a))
128+ -- ^ Current CSV parser state.
129+ -> InputStream ByteString
130+ -- ^ Upstream.
131+ -> IO (Maybe (Either String a))
132+ -- ^ Data feed downstream.
133+dispatch queueRef parserRef input = do
134+ queue <- readIORef queueRef
135+
136+ case queue of
137+ [] -> do
138+ parser <- readIORef parserRef
139+ case parser of
140+ Nothing -> return Nothing
141+ Just (Fail _ e) -> bomb ("input data malformed: " ++ e)
142+ Just (Many [] f) -> more f
143+ Just (Many xs f) -> writeIORef parserRef (Just $ Many [] f) >> feed xs
144+ Just (Done xs ) -> writeIORef parserRef Nothing >> feed xs
145+
146+ (x:xs) -> do
147+ writeIORef queueRef xs
148+ return (Just x)
149+
150+ where
151+ -- Send more data to the CSV parser. If there is no more data
152+ -- from upstream then send an empty @ByteString@.
153+ more :: (ByteString -> Parser a) -> IO (Maybe (Either String a))
154+ more f = do bs <- Streams.read input
155+ writeIORef parserRef (Just $ maybe (f BS.empty) f bs)
156+ dispatch queueRef parserRef input
157+
158+ -- Feed records downstream.
159+ feed :: [Either String a] -> IO (Maybe (Either String a))
160+ feed xs = do writeIORef queueRef xs
161+ dispatch queueRef parserRef input
162+
163+--------------------------------------------------------------------------------
164+-- | Throw an exception.
165+bomb :: String -> IO a
166+bomb = throwIO . StreamDecodingError