Some utilities for use with the pipes package

PipeUtil.hs

module PipeUtil where

import Prelude hiding (getContents)
import Control.Arrow
import Control.Monad
import Control.Proxy
import System.IO (openFile, hClose, hIsEOF, hIsClosed, Handle, IOMode(ReadMode))
import Control.Monad.Trans
import Control.Monad.Trans.Resource

import Codec.Text.IConv (convertStrictly, EncodingName, ConversionError(..))
import qualified Data.Text.Lazy as TL
import qualified Data.Text.Lazy.Encoding as TL
import qualified Data.ByteString.Lazy as LZ

-- | Take an IO action indicating when there is no more data to read,
-- an IO action that cleans up when reading is done, and a function to
-- read n more, and gives a Sever that responds to requests for n more
-- until there is no more.
getContents :: IO Bool -> IO () -> (Int -> IO a) -> Int -> Server Int a (ResourceT IO) ()
getContents ioDone allDone io count = lift (register allDone) >> go count
	where
	go count = do
		isDone <- lift $ lift ioDone
		unless isDone (lift (lift $ io count) >>= respond >>= go)

hGetContents :: Handle -> (Handle -> Int -> IO a) -> Int -> Server Int a (ResourceT IO) ()
hGetContents h io = getContents (hIsClosed h ||. hIsEOF h) (hClose h) (io h)
	where
	a ||. b = a >>= (\a' -> if a' then return True else b)

-- | 'getContents' of a file specified by 'FilePath'
--
-- Don't need to catch exceptions raised by openFile, since we do not need
-- to close in that case (we never opened).
readFile :: FilePath -> (Handle -> Int -> IO a) -> Int -> Server Int a (ResourceT IO) ()
readFile pth io i = do
	h <- lift $ lift (openFile pth ReadMode)
	hGetContents h io i

-- | Take some n and always request from the upstream server in chunks
-- of size n
getChunks :: (Monad m) => Int -> () -> Proxy Int a () a m r
getChunks count = foreverK (\() -> request count >>= respond)

data UnGet a = Get Int | UnGet a
data UnGot a = Got a   | UnGot

-- | Take a length function, a splitAt function, an an append function and
-- add the ability to "unget" to any Server for which such functions can be
-- defined
canUnGet :: (Monad m) =>
	(a -> Int)              -- ^ length
	-> (Int -> a -> (a, a)) -- ^ splitAt
	-> (a -> a -> a)        -- ^ append
	-> UnGet a -> Proxy Int a (UnGet a) (UnGot a) m r
canUnGet theLength theSplitAt theAppend = go []
	where
	getChunk [] 0 = error "Programmer error in canUnGet"
	getChunk [] count = liftM (flip (,) []) (request count)
	getChunk ((x,len):xs) count
		| count > len =
			first (x `theAppend`) `liftM` getChunk xs (count - len)
		| count <= len =
			let (use,rest) = theSplitAt count x in
				return (use, (rest,len - count):xs)
	getChunk _ _ = error "Programmer error in canUnGet"

	go buf (Get count) = do
		(r, buf') <- getChunk buf count
		respond (Got r) >>= go buf'
	go buf (UnGet x) = respond UnGot >>= go ((x, theLength x):buf)

-- | Try to decode a Lazy 'ByteString' upstream into 'Text' using 'iconv'
--
-- Has a bug: if you ask for more characters than are still in the file,
-- it terminates early (if the upstream terminates on EOF instead of
-- continuing to respond with blank ByteStrings).
decodeText :: (Monad m) => EncodingName -> Int -> Proxy Int LZ.ByteString Int TL.Text m r
decodeText inputEncoding = foreverK (\count ->
		respond =<< liftM TL.decodeUtf32BE
			(getUtf16 (fromIntegral count) =<< request count)
	)
	where
	getUtf16 count bytes = do
		case convertStrictly inputEncoding "UTF-32BE" bytes of
			Left s -- UTF-32BE is 4 bytes per Char
				| (LZ.length s `div` 4) >= count -> return s
				| otherwise -> do
					b <- request (fromIntegral $ count - (LZ.length s `div` 4))
					if LZ.null b then return s else
						getUtf16 count (bytes `LZ.append` b)
			Right (IncompleteChar _) -> do
				b <- request 1
				when (LZ.null b) (error "Bad encoding") -- TODO: handle bad encodings better
				getUtf16 count (bytes `LZ.append` b)
			Right _ -> error "Bad encoding" -- TODO: handle bad encodings better

-- | Simple debug printer that just drains the Session and calls 'print'
debugPrint :: (Show a, MonadIO m) => () -> Client () a m ()
debugPrint = foreverK debugPrint1

-- | Simple debug printer that prints a single item and exits
debugPrint1 :: (Show a, MonadIO m) => () -> Client () a m ()
debugPrint1 () = request () >>= lift . liftIO . print