module Data.Enumerator.Binary (
enumHandle
, enumFile
, iterHandle
, Data.Enumerator.Binary.head
, Data.Enumerator.Binary.drop
, Data.Enumerator.Binary.dropWhile
, Data.Enumerator.Binary.take
, Data.Enumerator.Binary.takeWhile
, Data.Enumerator.Binary.consume
, require
, isolate
) where
import Prelude hiding (head, drop, takeWhile)
import Data.Enumerator hiding (head, drop)
import qualified Data.ByteString as B
import Data.Enumerator.Util (tryStep)
import qualified Control.Exception as Exc
import Control.Monad.IO.Class (MonadIO)
import qualified System.IO as IO
import System.IO.Error (isEOFError)
import Data.Word (Word8)
import qualified Data.ByteString.Lazy as BL
enumHandle :: MonadIO m
=> Integer
-> IO.Handle
-> Enumerator B.ByteString m b
enumHandle bufferSize h = loop where
loop (Continue k) = withBytes $ \bytes ->
if B.null bytes
then continue k
else k (Chunks [bytes]) >>== loop
loop step = returnI step
intSize = fromInteger bufferSize
withBytes = tryStep $ do
hasInput <- Exc.catch
(IO.hWaitForInput h (1))
(\err -> if isEOFError err
then return False
else Exc.throwIO err)
if hasInput
then B.hGetNonBlocking h intSize
else return B.empty
enumFile :: FilePath -> Enumerator B.ByteString IO b
enumFile path = enum where
withHandle = tryStep (IO.openBinaryFile path IO.ReadMode)
enum step = withHandle $ \h -> do
Iteratee $ Exc.finally
(runIteratee (enumHandle 4096 h step))
(IO.hClose h)
iterHandle :: MonadIO m => IO.Handle
-> Iteratee B.ByteString m ()
iterHandle h = continue step where
step EOF = yield () EOF
step (Chunks []) = continue step
step (Chunks bytes) = let
put = mapM_ (B.hPut h) bytes
in tryStep put (\_ -> continue step)
toChunks :: BL.ByteString -> Stream B.ByteString
toChunks = Chunks . BL.toChunks
head :: Monad m => Iteratee B.ByteString m (Maybe Word8)
head = continue loop where
loop (Chunks xs) = case BL.uncons (BL.fromChunks xs) of
Just (char, extra) -> yield (Just char) (toChunks extra)
Nothing -> head
loop EOF = yield Nothing EOF
drop :: Monad m => Integer -> Iteratee B.ByteString m ()
drop n | n <= 0 = return ()
drop n = continue (loop n) where
loop n' (Chunks xs) = iter where
lazy = BL.fromChunks xs
len = toInteger (BL.length lazy)
iter = if len < n'
then drop (n' len)
else yield () (toChunks (BL.drop (fromInteger n') lazy))
loop _ EOF = yield () EOF
dropWhile :: Monad m => (Word8 -> Bool) -> Iteratee B.ByteString m ()
dropWhile p = continue loop where
loop (Chunks xs) = iter where
lazy = BL.dropWhile p (BL.fromChunks xs)
iter = if BL.null lazy
then continue loop
else yield () (toChunks lazy)
loop EOF = yield () EOF
take :: Monad m => Integer -> Iteratee B.ByteString m BL.ByteString
take n | n <= 0 = return BL.empty
take n = continue (loop id n) where
loop acc n' (Chunks xs) = iter where
lazy = BL.fromChunks xs
len = toInteger (BL.length lazy)
iter = if len < n'
then continue (loop (acc . (BL.append lazy)) (n' len))
else let
(xs', extra) = BL.splitAt (fromInteger n') lazy
in yield (acc xs') (toChunks extra)
loop acc _ EOF = yield (acc BL.empty) EOF
takeWhile :: Monad m => (Word8 -> Bool) -> Iteratee B.ByteString m BL.ByteString
takeWhile p = continue (loop id) where
loop acc (Chunks []) = continue (loop acc)
loop acc (Chunks xs) = iter where
lazy = BL.fromChunks xs
(xs', extra) = BL.span p lazy
iter = if BL.null extra
then continue (loop (acc . (BL.append lazy)))
else yield (acc xs') (toChunks extra)
loop acc EOF = yield (acc BL.empty) EOF
consume :: Monad m => Iteratee B.ByteString m BL.ByteString
consume = continue (loop id) where
loop acc (Chunks []) = continue (loop acc)
loop acc (Chunks xs) = iter where
lazy = BL.fromChunks xs
iter = continue (loop (acc . (BL.append lazy)))
loop acc EOF = yield (acc BL.empty) EOF
require :: Monad m => Integer -> Iteratee B.ByteString m ()
require n | n <= 0 = return ()
require n = continue (loop id n) where
loop acc n' (Chunks xs) = iter where
lazy = BL.fromChunks xs
len = toInteger (BL.length lazy)
iter = if len < n'
then continue (loop (acc . (BL.append lazy)) (n' len))
else yield () (toChunks (acc lazy))
loop _ _ EOF = throwError (Exc.ErrorCall "require: Unexpected EOF")
isolate :: Monad m => Integer -> Enumeratee B.ByteString B.ByteString m b
isolate n step | n <= 0 = return step
isolate n (Continue k) = continue loop where
loop (Chunks []) = continue loop
loop (Chunks xs) = iter where
lazy = BL.fromChunks xs
len = toInteger (BL.length lazy)
iter = if len <= n
then k (Chunks xs) >>== isolate (n len)
else let
(s1, s2) = BL.splitAt (fromInteger n) lazy
in k (toChunks s1) >>== (\step -> yield step (toChunks s2))
loop EOF = k EOF >>== (\step -> yield step EOF)
isolate n step = drop n >> return step