module Data.Conduit.Binary
( sourceFile
, sourceHandle
, sourceIOHandle
, sourceFileRange
, sinkFile
, sinkHandle
, sinkIOHandle
, conduitFile
, isolate
, openFile
, head
, takeWhile
, dropWhile
, take
, Data.Conduit.Binary.lines
) where
import Prelude hiding (head, take, takeWhile, dropWhile)
import qualified Data.ByteString as S
import qualified Data.ByteString.Lazy as L
import Data.Conduit
import qualified Data.Conduit.List as CL
import Control.Exception (assert)
import Control.Monad (liftM)
import Control.Monad.IO.Class (liftIO)
import qualified System.IO as IO
import Control.Monad.Trans.Resource (allocate, release)
import Data.Word (Word8)
import Data.Monoid (mempty)
#if CABAL_OS_WINDOWS
import qualified System.Win32File as F
#elif NO_HANDLES
import qualified System.PosixFile as F
#endif
openFile :: MonadResource m
=> FilePath
-> IO.IOMode
-> m IO.Handle
openFile fp mode = liftM snd $ allocate (IO.openBinaryFile fp mode) IO.hClose
sourceFile :: MonadResource m
=> FilePath
-> Source m S.ByteString
sourceFile fp =
#if CABAL_OS_WINDOWS || NO_HANDLES
sourceIO (F.openRead fp)
F.close
(liftIO . F.read)
#else
sourceIOHandle (IO.openBinaryFile fp IO.ReadMode)
#endif
sourceHandle :: MonadResource m
=> IO.Handle
-> Source m S.ByteString
sourceHandle h =
src
where
src = SourceM pull close
pull = do
bs <- liftIO (S.hGetSome h 4096)
if S.null bs
then return Closed
else return $ Open src close bs
close = return ()
sourceIOHandle :: MonadResource m
=> IO IO.Handle
-> Source m S.ByteString
sourceIOHandle alloc = sourceIO alloc IO.hClose
(\handle -> do
bs <- liftIO (S.hGetSome handle 4096)
if S.null bs
then return IOClosed
else return $ IOOpen bs)
sinkHandle :: MonadResource m
=> IO.Handle
-> Sink S.ByteString m ()
sinkHandle h =
Processing push close
where
push input = SinkM $ liftIO (S.hPut h input) >> return (Processing push close)
close = return ()
sinkIOHandle :: MonadResource m
=> IO IO.Handle
-> Sink S.ByteString m ()
sinkIOHandle alloc = sinkIO alloc IO.hClose
(\handle bs -> liftIO (S.hPut handle bs) >> return IOProcessing)
(const $ return ())
sourceFileRange :: MonadResource m
=> FilePath
-> Maybe Integer
-> Maybe Integer
-> Source m S.ByteString
sourceFileRange fp offset count = SourceM
(do
(key, handle) <- allocate (IO.openBinaryFile fp IO.ReadMode) IO.hClose
case offset of
Nothing -> return ()
Just off -> liftIO $ IO.hSeek handle IO.AbsoluteSeek off
case count of
Nothing -> pullUnlimited handle key
Just c -> pullLimited c handle key)
(return ())
where
pullUnlimited handle key = do
bs <- liftIO $ S.hGetSome handle 4096
if S.null bs
then do
release key
return Closed
else do
let src = SourceM
(pullUnlimited handle key)
(release key)
return $ Open src (release key) bs
pullLimited c0 handle key = do
let c = fromInteger c0
bs <- liftIO $ S.hGetSome handle (min c 4096)
let c' = c S.length bs
assert (c' >= 0) $
if S.null bs
then do
release key
return Closed
else do
let src = SourceM
(pullLimited (toInteger c') handle key)
(release key)
return $ Open src (release key) bs
sinkFile :: MonadResource m
=> FilePath
-> Sink S.ByteString m ()
sinkFile fp = sinkIOHandle (IO.openBinaryFile fp IO.WriteMode)
conduitFile :: MonadResource m
=> FilePath
-> Conduit S.ByteString m S.ByteString
conduitFile fp = conduitIO
(IO.openBinaryFile fp IO.WriteMode)
IO.hClose
(\handle bs -> do
liftIO $ S.hPut handle bs
return $ IOProducing [bs])
(const $ return [])
isolate :: Monad m
=> Int
-> Conduit S.ByteString m S.ByteString
isolate count0 = conduitState
count0
push
close
where
push 0 bs = return $ StateFinished (Just bs) []
push count bs = do
let (a, b) = S.splitAt count bs
let count' = count S.length a
return $
if count' == 0
then StateFinished (if S.null b then Nothing else Just b) (if S.null a then [] else [a])
else assert (S.null b) $ StateProducing count' [a]
close _ = return []
head :: Monad m => Sink S.ByteString m (Maybe Word8)
head =
Processing push close
where
push bs =
case S.uncons bs of
Nothing -> Processing push close
Just (w, bs') ->
let lo = if S.null bs' then Nothing else Just bs'
in Done lo (Just w)
close = return Nothing
takeWhile :: Monad m => (Word8 -> Bool) -> Conduit S.ByteString m S.ByteString
takeWhile p =
NeedInput push close
where
push bs
| S.null y =
let r = NeedInput push close
in if S.null x
then r
else HaveOutput r (return ()) x
| otherwise =
let f = Finished $ Just y
in if S.null x
then f
else HaveOutput f (return ()) x
where
(x, y) = S.span p bs
close = mempty
dropWhile :: Monad m => (Word8 -> Bool) -> Sink S.ByteString m ()
dropWhile p =
Processing push close
where
push bs
| S.null bs' = Processing push close
| otherwise = Done (Just bs') ()
where
bs' = S.dropWhile p bs
close = return ()
take :: Monad m => Int -> Sink S.ByteString m L.ByteString
take n = L.fromChunks `liftM` (isolate n =$ CL.consume)
lines :: Monad m => Conduit S.ByteString m S.ByteString
lines =
conduitState id push close
where
push front bs' = return $ StateProducing leftover ls
where
bs = front bs'
(leftover, ls) = getLines id bs
getLines front bs
| S.null bs = (id, front [])
| S.null y = (S.append x, front [])
| otherwise = getLines (front . (x:)) (S.drop 1 y)
where
(x, y) = S.breakByte 10 bs
close front
| S.null bs = return []
| otherwise = return [bs]
where
bs = front S.empty