@@ -0,0 +1,120 @@
1+{-
2+
3+This file is part of the Haskell package cassava-streams. It is
4+subject to the license terms in the LICENSE file found in the
5+top-level directory of this distribution and at
6+git://pmade.com/cassava-streams/LICENSE. No part of cassava-streams
7+package, including this file, may be copied, modified, propagated, or
8+distributed except according to the terms contained in the LICENSE
9+file.
10+
11+-}
12+
13+--------------------------------------------------------------------------------
14+module System.IO.Streams.Csv.Decode
15+ ( decodeStream
16+ , decodeStreamWith
17+ , decodeStreamByName
18+ , decodeStreamByNameWith
19+ ) where
20+
21+--------------------------------------------------------------------------------
22+import Data.ByteString (ByteString)
23+import qualified Data.ByteString as BS
24+import Data.Csv hiding (Parser, decodeWith, decodeByNameWith)
25+import Data.Csv.Incremental
26+import Data.Either (lefts, rights)
27+import Data.IORef
28+import System.IO.Streams (InputStream, makeInputStream)
29+import qualified System.IO.Streams as Streams
30+
31+--------------------------------------------------------------------------------
32+-- | Create an @InputStream@ which decodes CSV records from the given
33+-- upstream data source.
34+--
35+-- Equivalent to @decodeStreamWith defaultDecodeOptions@.
36+decodeStream :: (FromRecord a)
37+ => HasHeader -- ^ Whether to skip a header or not.
38+ -> InputStream ByteString -- ^ Upstream.
39+ -> IO (InputStream a) -- ^ An @InputStream@ which produces records.
40+decodeStream = decodeStreamWith defaultDecodeOptions
41+
42+--------------------------------------------------------------------------------
43+-- | Create an @InputStream@ which decodes CSV records from the given
44+-- upstream data source.
45+decodeStreamWith :: (FromRecord a)
46+ => DecodeOptions -- ^ CSV decoding options.
47+ -> HasHeader -- ^ Whether to skip a header or not.
48+ -> InputStream ByteString -- ^ Upstream.
49+ -> IO (InputStream a) -- ^ An @InputStream@ which produces records.
50+decodeStreamWith ops hdr input = do
51+ queue <- newIORef []
52+ parser <- newIORef (decodeWith ops hdr)
53+ makeInputStream (dispatch queue parser input)
54+
55+--------------------------------------------------------------------------------
56+-- | Create an @InputStream@ which decodes CSV records from the given
57+-- upstream data source. Data should be preceded by a header.
58+--
59+-- Equivalent to @decodeStreamByNameWith defaultDecodeOptions@.
60+decodeStreamByName :: (FromNamedRecord a)
61+ => InputStream ByteString -- ^ Upstream.
62+ -> IO (InputStream a) -- ^ An @InputStream@ which produces records.
63+decodeStreamByName = decodeStreamByNameWith defaultDecodeOptions
64+
65+--------------------------------------------------------------------------------
66+-- | Create an @InputStream@ which decodes CSV records from the given
67+-- upstream data source. Data should be preceded by a header.
68+decodeStreamByNameWith :: (FromNamedRecord a)
69+ => DecodeOptions -- ^ CSV decoding options.
70+ -> InputStream ByteString -- ^ Upstream.
71+ -> IO (InputStream a) -- ^ An @InputStream@ which produces records.
72+decodeStreamByNameWith ops input = go (decodeByNameWith ops) where
73+ -- Dispatch on the HeaderParser type.
74+ go (FailH _ e) = bomb e
75+ go (PartialH f) = Streams.read input >>= go . maybe (f BS.empty) f
76+ go (DoneH _ p) = do
77+ queue <- newIORef []
78+ parser <- newIORef p
79+ makeInputStream (dispatch queue parser input)
80+
81+--------------------------------------------------------------------------------
82+-- | Internal function which feeds data to the CSV parser.
83+dispatch :: IORef [a] -- ^ List of queued CSV records.
84+ -> IORef (Parser a) -- ^ Current CSV parser state.
85+ -> InputStream ByteString -- ^ Upstream.
86+ -> IO (Maybe a) -- ^ Data feed downstream.
87+dispatch queueRef parserRef input = do
88+ queue <- readIORef queueRef
89+
90+ case queue of
91+ [] -> do
92+ parser <- readIORef parserRef
93+ case parser of
94+ Fail _ e -> bomb e
95+ Many xs f -> more f >> feed xs
96+ Done xs -> feed xs
97+
98+ (x:xs) -> do
99+ writeIORef queueRef xs
100+ return (Just x)
101+
102+ where
103+ -- Send more data to the CSV parser. If there is no more data
104+ -- from upstream then send an empty @ByteString@.
105+ more f = do
106+ bytes <- Streams.read input
107+ writeIORef parserRef (maybe (f BS.empty) f bytes)
108+
109+ -- Feed more data downstream or fail if some records didn't parse
110+ -- correctly. The elements are wrapped in an @Either@ which
111+ -- indicates decoding failures.
112+ feed xs = if not $ null (lefts xs)
113+ then bomb $ concat (lefts xs)
114+ else writeIORef queueRef (rights xs) >>
115+ dispatch queueRef parserRef input
116+
117+--------------------------------------------------------------------------------
118+-- | FIXME: use a proper exception.
119+bomb :: String -> IO a
120+bomb = fail