module Data.Conduit.Binary
(
sourceFile
, sourceHandle
, sourceIOHandle
, sourceFileRange
, sinkFile
, sinkHandle
, sinkIOHandle
, conduitFile
, sourceLbs
, head
, dropWhile
, take
, drop
, isolate
, takeWhile
, Data.Conduit.Binary.lines
) where
import Prelude hiding (head, take, drop, takeWhile, dropWhile)
import qualified Data.ByteString as S
import qualified Data.ByteString.Lazy as L
import Data.Conduit
import Data.Conduit.List (sourceList)
import Control.Exception (assert)
import Control.Monad (unless)
import Control.Monad.IO.Class (liftIO, MonadIO)
import qualified System.IO as IO
import Data.Word (Word8)
import Control.Applicative ((<$>))
#if CABAL_OS_WINDOWS
import qualified System.Win32File as F
#elif NO_HANDLES
import qualified System.PosixFile as F
#endif
sourceFile :: MonadResource m
=> FilePath
-> Producer m S.ByteString
sourceFile fp =
#if CABAL_OS_WINDOWS || NO_HANDLES
bracketP
(F.openRead fp)
F.close
loop
where
loop h = liftIO (F.read h) >>= maybe (return ()) (\bs -> yield bs >> loop h)
#else
sourceIOHandle (IO.openBinaryFile fp IO.ReadMode)
#endif
sourceHandle :: MonadIO m
=> IO.Handle
-> Producer m S.ByteString
sourceHandle h =
loop
where
loop = do
bs <- liftIO (S.hGetSome h 4096)
if S.null bs
then return ()
else yield bs >> loop
sourceIOHandle :: MonadResource m
=> IO IO.Handle
-> Producer m S.ByteString
sourceIOHandle alloc = bracketP alloc IO.hClose sourceHandle
sinkHandle :: MonadIO m
=> IO.Handle
-> Consumer S.ByteString m ()
sinkHandle h = awaitForever $ liftIO . S.hPut h
sinkIOHandle :: MonadResource m
=> IO IO.Handle
-> Consumer S.ByteString m ()
sinkIOHandle alloc = bracketP alloc IO.hClose sinkHandle
sourceFileRange :: MonadResource m
=> FilePath
-> Maybe Integer
-> Maybe Integer
-> Producer m S.ByteString
sourceFileRange fp offset count = bracketP
(IO.openBinaryFile fp IO.ReadMode)
IO.hClose
start
where
start handle = do
case offset of
Nothing -> return ()
Just off -> liftIO $ IO.hSeek handle IO.AbsoluteSeek off
case count of
Nothing -> pullUnlimited handle
Just c -> pullLimited (fromInteger c) handle
pullUnlimited handle = do
bs <- liftIO $ S.hGetSome handle 4096
if S.null bs
then return ()
else do
yield bs
pullUnlimited handle
pullLimited c handle = do
bs <- liftIO $ S.hGetSome handle (min c 4096)
let c' = c S.length bs
assert (c' >= 0) $
if S.null bs
then return ()
else do
yield bs
pullLimited c' handle
sinkFile :: MonadResource m
=> FilePath
-> Consumer S.ByteString m ()
sinkFile fp = sinkIOHandle (IO.openBinaryFile fp IO.WriteMode)
conduitFile :: MonadResource m
=> FilePath
-> Conduit S.ByteString m S.ByteString
conduitFile fp = bracketP
(IO.openBinaryFile fp IO.WriteMode)
IO.hClose
go
where
go h = awaitForever $ \bs -> liftIO (S.hPut h bs) >> yield bs
isolate :: Monad m
=> Int
-> Conduit S.ByteString m S.ByteString
isolate =
loop
where
loop 0 = return ()
loop count = do
mbs <- await
case mbs of
Nothing -> return ()
Just bs -> do
let (a, b) = S.splitAt count bs
case count S.length a of
0 -> do
unless (S.null b) $ leftover b
yield a
count' -> assert (S.null b) $ yield a >> loop count'
head :: Monad m => Consumer S.ByteString m (Maybe Word8)
head = do
mbs <- await
case mbs of
Nothing -> return Nothing
Just bs ->
case S.uncons bs of
Nothing -> head
Just (w, bs') -> leftover bs' >> return (Just w)
takeWhile :: Monad m => (Word8 -> Bool) -> Conduit S.ByteString m S.ByteString
takeWhile p =
loop
where
loop = await >>= maybe (return ()) go
go bs
| S.null x = next
| otherwise = yield x >> next
where
next = if S.null y then loop else leftover y
(x, y) = S.span p bs
dropWhile :: Monad m => (Word8 -> Bool) -> Consumer S.ByteString m ()
dropWhile p =
loop
where
loop = do
mbs <- await
case S.dropWhile p <$> mbs of
Nothing -> return ()
Just bs
| S.null bs -> loop
| otherwise -> leftover bs
take :: Monad m => Int -> Consumer S.ByteString m L.ByteString
take n0 =
go n0 id
where
go n front =
await >>= maybe (return $ L.fromChunks $ front []) go'
where
go' bs =
case S.length bs `compare` n of
LT -> go (n S.length bs) (front . (bs:))
EQ -> return $ L.fromChunks $ front [bs]
GT ->
let (x, y) = S.splitAt n bs
in assert (not $ S.null y) $ leftover y >> return (L.fromChunks $ front [x])
drop :: Monad m => Int -> Consumer S.ByteString m ()
drop =
go
where
go n =
await >>= maybe (return ()) go'
where
go' bs =
case S.length bs `compare` n of
LT -> go (n S.length bs)
EQ -> return ()
GT ->
let y = S.drop n bs
in assert (not $ S.null y) $ leftover y >> return ()
lines :: Monad m => Conduit S.ByteString m S.ByteString
lines =
loop id
where
loop front = await >>= maybe (finish front) (go front)
finish front =
let final = front S.empty
in unless (S.null final) (yield final)
go sofar more =
case S.uncons second of
Just (_, second') -> yield (sofar first) >> go id second'
Nothing ->
let rest = sofar more
in loop $ S.append rest
where
(first, second) = S.breakByte 10 more
sourceLbs :: Monad m => L.ByteString -> Producer m S.ByteString
sourceLbs = sourceList . L.toChunks