module Network.Wai.Enumerator
(
mapE
,
toLBS
, fromLBS
, fromLBS'
, toSource
, fromHandle
, fromFile
, fromEitherFile
, buffer
) where
import Network.Wai (Enumerator (..), Source (..))
import qualified Network.Wai.Source as Source
import qualified Data.ByteString.Lazy as L
import qualified Data.ByteString as B
import System.IO (withBinaryFile, IOMode (ReadMode), Handle, hIsEOF)
import Data.ByteString.Lazy.Internal (defaultChunkSize)
import Control.Concurrent (forkIO)
import Control.Concurrent.MVar
import Control.Monad ((<=<))
mapE :: (B.ByteString -> B.ByteString) -> Enumerator -> Enumerator
mapE f (Enumerator e) = Enumerator $ \iter -> e (iter' iter) where
iter' iter a = iter a . f
toLBS :: Enumerator -> IO L.ByteString
toLBS = Source.toLBS <=< toSource
fromLBS :: L.ByteString -> Enumerator
fromLBS lbs = Enumerator $ \iter a0 -> helper iter a0 $ L.toChunks lbs where
helper _ a [] = return $ Right a
helper iter a (x:xs) = do
ea <- iter a x
case ea of
Left a' -> return $ Left a'
Right a' -> helper iter a' xs
fromLBS' :: IO L.ByteString -> Enumerator
fromLBS' lbs' = Enumerator $ \iter a0 -> lbs' >>= \lbs ->
runEnumerator (fromLBS lbs) iter a0
toSource :: Enumerator -> IO Source
toSource (Enumerator e) = do
buff <- newEmptyMVar
_ <- forkIO $ e (helper buff) () >> putMVar buff Nothing
return $ source buff
where
helper :: MVar (Maybe B.ByteString)
-> ()
-> B.ByteString
-> IO (Either () ())
helper buff _ bs = do
putMVar buff $ Just bs
return $ Right ()
source :: MVar (Maybe B.ByteString)
-> Source
source mmbs = Source $ do
mbs <- takeMVar mmbs
case mbs of
Nothing -> do
putMVar mmbs Nothing
return Nothing
Just bs -> return $ Just (bs, source mmbs)
fromHandle :: Handle -> Enumerator
fromHandle h = Enumerator $ \iter a -> do
eof <- hIsEOF h
if eof
then return $ Right a
else do
bs <- B.hGet h defaultChunkSize
ea' <- iter a bs
case ea' of
Left a' -> return $ Left a'
Right a' -> runEnumerator (fromHandle h) iter a'
fromFile :: FilePath -> Enumerator
fromFile fp = Enumerator $ \iter a0 -> withBinaryFile fp ReadMode $ \h ->
runEnumerator (fromHandle h) iter a0
fromEitherFile :: Either FilePath Enumerator -> Enumerator
fromEitherFile = either fromFile id
buffer :: Enumerator -> Enumerator
buffer (Enumerator e) =
Enumerator go
where
go iter seed = do
res <- e (iter' iter) (B.empty, seed)
case res of
Left (_, seed') -> return $ Left seed'
Right (buff, seed') -> iter seed' buff
iter' iter (buff, seed) bs = do
if B.length buff + B.length bs < defaultChunkSize
then return $ Right (buff `B.append` bs, seed)
else do
let (x, y) = B.splitAt (defaultChunkSize B.length buff) bs
res <- iter seed $ buff `B.append` x
case res of
Left seed' -> return $ Left (y, seed')
Right seed' -> return $ Right (y, seed')