IO.hs

-----------------------------------------------------------------------------
-- |
-- Module: Data.Enumerator.IO
-- Copyright: 2010 John Millikin
-- License: MIT
--
-- Maintainer: jmillikin@gmail.com
-- Portability: portable
--
-- Enumerator-based IO
--
-----------------------------------------------------------------------------
module Data.Enumerator.IO
	( enumHandle
	, enumFile
	, iterFile
	, iterHandle
	) where
import Data.Enumerator
import Control.Monad.IO.Class (liftIO)
import qualified Control.Exception as E
import qualified Data.ByteString as B
import qualified Foreign as F
import qualified System.IO as IO
-- | Read bytes (in chunks of the given buffer size) from the handle, and
-- stream them to an 'Iteratee'. If an exception occurs during file IO,
-- enumeration will stop and 'Error' will be returned. Exceptions from the
-- iteratee are not caught.
enumHandle :: Integer -- ^ Buffer size
           -> IO.Handle
           -> Enumerator E.SomeException B.ByteString IO b
enumHandle bufferSize h = Iteratee . F.allocaBytes size' . loop where
	size' = fromInteger bufferSize
	loop (Continue k) = read' k
	loop step = const $ return step
	read' k p = do
		eitherN <- E.try $ IO.hGetBuf h p size'
		case eitherN of
			Left err -> return $ Error err
			Right 0 -> return $ Continue k
			Right n -> do
				bytes <- B.packCStringLen (p, n)
				step <- runIteratee (k (Chunks [bytes]))
				loop step p
-- | Opens a file path in binary mode, and passes the handle to 'enumHandle'.
-- The file will be closed when the 'Iteratee' finishes.
enumFile :: FilePath -> Enumerator E.SomeException B.ByteString IO b
enumFile path s = Iteratee $ do
	eitherH <- E.try $ IO.openBinaryFile path IO.ReadMode
	case eitherH of
		Left err -> return $ Error err
		Right h -> E.finally
			(runIteratee (enumHandle 4096 h s))
			(IO.hClose h)
-- | Read bytes from a stream and write them to a handle. If an exception
-- occurs during file IO, enumeration will stop and 'Error' will be
-- returned.
iterHandle :: IO.Handle -> Iteratee E.SomeException B.ByteString IO ()
iterHandle h = continue step where
	step EOF = yield () EOF
	step (Chunks bytes) = do
		eitherErr <- liftIO . E.try $ mapM_ (B.hPut h) bytes
		case eitherErr of
			Left err -> throwError err
			_ -> continue step
-- | Opens a file path in binary mode, and passes the handle to 'iterHandle'.
-- The file will be closed when the 'Iteratee' finishes.
iterFile :: FilePath -> Iteratee E.SomeException B.ByteString IO ()
iterFile path = Iteratee $ do
	eitherH <- E.try $ IO.openBinaryFile path IO.WriteMode
	case eitherH of
		Left err -> return $ Error err
		Right h -> E.finally
			(runIteratee (iterHandle h))
			(IO.hClose h)