module Text.JSON.YAJL.Enumerator
( Event (..)
, parseBytesIO
, parseTextIO
, parseBytesST
, parseTextST
, Y.GeneratorConfig (..)
, Y.GeneratorError (..)
, generateBytesIO
, generateTextIO
, generateBytesST
, generateTextST
) where
import Prelude hiding (null)
import qualified Prelude as Prelude
import qualified Data.ByteString as B
import qualified Data.Text as T
import qualified Data.Text.Encoding as TE
import qualified Data.Enumerator as E
import Data.Enumerator ((>>==))
import qualified Text.JSON.YAJL as Y
import Control.Exception as Exc
import Control.Monad.IO.Class (MonadIO, liftIO)
import Control.Monad.Trans.Class (lift)
import Control.Monad.ST (ST, stToIO, unsafeIOToST, unsafeSTToIO, RealWorld)
import qualified Data.STRef as ST
import qualified Data.IORef as IO
data Event
= EventNull
| EventBoolean Bool
| EventNumber B.ByteString
| EventText T.Text
| EventBeginArray
| EventEndArray
| EventBeginObject
| EventAttributeName T.Text
| EventEndObject
deriving (Show)
parseBytesIO :: MonadIO m => E.Enumeratee B.ByteString Event m b
parseBytesIO = parseIO Y.parseBytes
parseTextIO :: MonadIO m => E.Enumeratee T.Text Event m b
parseTextIO = parseIO Y.parseText
parseBytesST :: E.Enumeratee B.ByteString Event (ST s) b
parseBytesST = parseST Y.parseBytes
parseTextST :: E.Enumeratee T.Text Event (ST s) b
parseTextST = parseST Y.parseText
newParserIO :: IO (Y.Parser IO, IO.IORef [Event])
newParserIO = do
p <- Y.newParserIO
eventRef <- IO.newIORef []
let addEvent e = IO.modifyIORef eventRef (e:) >> return True
setCallbacks p addEvent
return (p, eventRef)
newParserST :: ST s (Y.Parser (ST s), ST.STRef s [Event])
newParserST = do
p <- Y.newParserST
eventRef <- ST.newSTRef []
let addEvent e = ST.modifySTRef eventRef (e:) >> return True
setCallbacks p addEvent
return (p, eventRef)
setCallbacks :: Monad m => Y.Parser m -> (Event -> m Bool) -> m ()
setCallbacks p addEvent = do
let set = Y.setCallback p
set Y.parsedBeginArray (addEvent EventBeginArray)
set Y.parsedEndArray (addEvent EventEndArray)
set Y.parsedBeginObject (addEvent EventBeginObject)
set Y.parsedEndObject (addEvent EventEndObject)
set Y.parsedNull (addEvent EventNull)
set Y.parsedBoolean (addEvent . EventBoolean)
set Y.parsedNumber (addEvent . EventNumber)
set Y.parsedAttributeText (addEvent . EventAttributeName)
set Y.parsedText (addEvent . EventText)
parseIO :: MonadIO m
=> (Y.Parser IO -> a -> IO Y.ParseStatus)
-> E.Enumeratee a Event m b
parseIO parseFn s = do
(p, eventRef) <- liftIO newParserIO
let withEvents io = liftIO $ do
IO.writeIORef eventRef []
status <- io
events <- IO.readIORef eventRef
return (reverse events, status)
let parseChunk bytes = withEvents (parseFn p bytes)
let complete = withEvents (Y.parseComplete p)
eneeParser parseChunk complete s
parseST :: (Y.Parser (ST s) -> a -> ST s Y.ParseStatus)
-> E.Enumeratee a Event (ST s) b
parseST parseFn s = do
(p, eventRef) <- lift newParserST
let withEvents st = do
ST.writeSTRef eventRef []
status <- st
events <- ST.readSTRef eventRef
return (reverse events, status)
let parseChunk bytes = withEvents (parseFn p bytes)
let complete = withEvents (Y.parseComplete p)
eneeParser parseChunk complete s
eneeParser :: Monad m
=> (a -> m ([Event], Y.ParseStatus))
-> m ([Event], Y.ParseStatus)
-> E.Enumeratee a Event m b
eneeParser parseChunk parseComplete = E.checkDone (E.continue . step) where
step k (E.Chunks xs) = parseLoop k xs
step k E.EOF = checkEvents k parseComplete
(\k' -> E.yield (E.Continue k') E.EOF)
(\_ -> throwError (T.pack "Unexpected EOF"))
parseLoop k [] = E.continue (step k)
parseLoop k (x:xs) = checkEvents k (parseChunk x)
(\k' -> E.yield (E.Continue k') (E.Chunks xs))
(\k' -> parseLoop k' xs)
checkEvents k getEvents onFinished onContinue = do
(events, status) <- lift getEvents
let checkError k' = case status of
Y.ParseError err -> throwError err
Y.ParseFinished -> onFinished k'
Y.ParseContinue -> onContinue k'
Y.ParseCancelled -> throwError (T.pack "Parse cancelled")
if null events
then checkError k
else k (E.Chunks events) >>== E.checkDone checkError
throwError = E.throwError . Exc.ErrorCall . T.unpack
class Nullable a where
null :: a -> Bool
instance Nullable [a] where
null = Prelude.null
instance Nullable B.ByteString where
null = B.null
instance Nullable T.Text where
null = T.null
generateTextIO :: MonadIO m
=> Y.GeneratorConfig
-> E.Enumeratee Event T.Text m b
generateTextIO = generateIO (fmap TE.decodeUtf8 . Y.getBuffer)
generateBytesIO :: MonadIO m
=> Y.GeneratorConfig
-> E.Enumeratee Event B.ByteString m b
generateBytesIO = generateIO Y.getBuffer
generateTextST :: Y.GeneratorConfig
-> E.Enumeratee Event T.Text (ST s) b
generateTextST = generateST (fmap TE.decodeUtf8 . Y.getBuffer)
generateBytesST :: Y.GeneratorConfig
-> E.Enumeratee Event B.ByteString (ST s) b
generateBytesST = generateST Y.getBuffer
generateIO :: (Nullable a, MonadIO m)
=> (Y.Generator RealWorld -> ST RealWorld a)
-> Y.GeneratorConfig
-> E.Enumeratee Event a m b
generateIO getBuf config s = do
g <- liftIO $ stToIO $ Y.newGenerator config
let takeBuf = liftIO $ stToIO $ do
buf <- getBuf g
Y.clearBuffer g
return buf
let genEvent e = liftIO $ Exc.handle (return . Just) $ do
stToIO $ genEventImpl g e
return Nothing
eneeGenerator genEvent takeBuf s
generateST :: Nullable a
=> (Y.Generator s -> ST s a)
-> Y.GeneratorConfig
-> E.Enumeratee Event a (ST s) b
generateST getBuf config s = do
g <- lift $ Y.newGenerator config
let takeBuf = do
buf <- getBuf g
Y.clearBuffer g
return buf
let genEvent e = unsafeIOToST $ Exc.handle (return . Just) $ do
unsafeSTToIO $ genEventImpl g e
return Nothing
eneeGenerator genEvent takeBuf s
eneeGenerator :: (Nullable a, Monad m)
=> (Event -> m (Maybe Y.GeneratorError))
-> m a
-> E.Enumeratee Event a m b
eneeGenerator genEvent takeBuf = E.checkDone (E.continue . step) where
step k (E.Chunks []) = E.continue (step k)
step k (E.Chunks xs) = parseLoop k xs
step k E.EOF = E.yield (E.Continue k) E.EOF
parseLoop k [] = checkBuf k (E.continue . step)
parseLoop k (x:xs) = do
maybeError <- lift $ genEvent x
case maybeError of
Just Y.GenerationComplete -> checkBuf k
(\k' -> E.yield (E.Continue k') (E.Chunks (x:xs)))
Just err -> checkBuf k (\_ -> E.throwError err)
Nothing -> parseLoop k xs
checkBuf k next = do
buf <- lift takeBuf
if null buf
then next k
else k (E.Chunks [buf]) >>== E.checkDone next
genEventImpl :: Y.Generator s -> Event -> ST s ()
genEventImpl g e = case e of
EventNull -> Y.generateNull g
EventBoolean x -> Y.generateBoolean g x
EventNumber num -> Y.generateNumber g num
EventText text -> Y.generateText g text
EventBeginArray -> Y.generateBeginArray g
EventEndArray -> Y.generateEndArray g
EventBeginObject -> Y.generateBeginObject g
EventAttributeName name -> Y.generateText g name
EventEndObject -> Y.generateEndObject g