module Data.Enumerator.Text (
enumHandle
, enumFile
, iterHandle
, Data.Enumerator.Text.head
, Data.Enumerator.Text.drop
, Data.Enumerator.Text.dropWhile
, Data.Enumerator.Text.take
, Data.Enumerator.Text.takeWhile
, Data.Enumerator.Text.consume
, require
, isolate
, Codec
, encode
, decode
, utf8
, utf16_le
, utf16_be
, utf32_le
, utf32_be
, ascii
, iso8859_1
) where
import qualified Prelude
import Prelude hiding (head, drop, takeWhile)
import Data.Enumerator hiding (head, drop)
import qualified Data.Text as T
import Data.Enumerator.Util (tryStep)
import qualified Data.Text.IO as TIO
import qualified Control.Exception as Exc
import Control.Monad.IO.Class (MonadIO)
import qualified System.IO as IO
import System.IO.Error (isEOFError)
import qualified Data.Text.Lazy as TL
import qualified Data.ByteString as B
import Data.Enumerator.Util (tSpanBy, tlSpanBy, reprWord, reprChar)
import Control.Arrow (first)
import Data.Bits ((.&.), (.|.), shiftL)
import Data.Char (ord)
import Data.Word (Word8, Word16)
import qualified Data.ByteString.Char8 as B8
import qualified Data.Text.Encoding as TE
import Data.Maybe (catMaybes)
import System.IO.Unsafe (unsafePerformIO)
enumHandle :: MonadIO m => IO.Handle
-> Enumerator T.Text m b
enumHandle h = loop where
loop (Continue k) = withText $ \maybeText ->
case maybeText of
Nothing -> continue k
Just text -> k (Chunks [text]) >>== loop
loop step = returnI step
withText = tryStep $ Exc.catch
(Just `fmap` TIO.hGetLine h)
(\err -> if isEOFError err
then return Nothing
else Exc.throwIO err)
enumFile :: FilePath -> Enumerator T.Text IO b
enumFile path = enum where
withHandle = tryStep (IO.openFile path IO.ReadMode)
enum step = withHandle $ \h -> Iteratee $ Exc.finally
(runIteratee (enumHandle h step))
(IO.hClose h)
iterHandle :: MonadIO m => IO.Handle
-> Iteratee T.Text m ()
iterHandle h = continue step where
step EOF = yield () EOF
step (Chunks []) = continue step
step (Chunks chunks) = let
put = mapM_ (TIO.hPutStr h) chunks
in tryStep put (\_ -> continue step)
toChunks :: TL.Text -> Stream T.Text
toChunks = Chunks . TL.toChunks
head :: Monad m => Iteratee T.Text m (Maybe Char)
head = continue loop where
loop (Chunks xs) = case TL.uncons (TL.fromChunks xs) of
Just (char, extra) -> yield (Just char) (toChunks extra)
Nothing -> head
loop EOF = yield Nothing EOF
drop :: Monad m => Integer -> Iteratee T.Text m ()
drop n | n <= 0 = return ()
drop n = continue (loop n) where
loop n' (Chunks xs) = iter where
lazy = TL.fromChunks xs
len = toInteger (TL.length lazy)
iter = if len < n'
then drop (n' len)
else yield () (toChunks (TL.drop (fromInteger n') lazy))
loop _ EOF = yield () EOF
dropWhile :: Monad m => (Char -> Bool) -> Iteratee T.Text m ()
dropWhile p = continue loop where
loop (Chunks xs) = iter where
lazy = TL.dropWhile p (TL.fromChunks xs)
iter = if TL.null lazy
then continue loop
else yield () (toChunks lazy)
loop EOF = yield () EOF
take :: Monad m => Integer -> Iteratee T.Text m TL.Text
take n | n <= 0 = return TL.empty
take n = continue (loop id n) where
loop acc n' (Chunks xs) = iter where
lazy = TL.fromChunks xs
len = toInteger (TL.length lazy)
iter = if len < n'
then continue (loop (acc . (TL.append lazy)) (n' len))
else let
(xs', extra) = TL.splitAt (fromInteger n') lazy
in yield (acc xs') (toChunks extra)
loop acc _ EOF = yield (acc TL.empty) EOF
takeWhile :: Monad m => (Char -> Bool) -> Iteratee T.Text m TL.Text
takeWhile p = continue (loop id) where
loop acc (Chunks []) = continue (loop acc)
loop acc (Chunks xs) = iter where
lazy = TL.fromChunks xs
(xs', extra) = tlSpanBy p lazy
iter = if TL.null extra
then continue (loop (acc . (TL.append lazy)))
else yield (acc xs') (toChunks extra)
loop acc EOF = yield (acc TL.empty) EOF
consume :: Monad m => Iteratee T.Text m TL.Text
consume = continue (loop id) where
loop acc (Chunks []) = continue (loop acc)
loop acc (Chunks xs) = iter where
lazy = TL.fromChunks xs
iter = continue (loop (acc . (TL.append lazy)))
loop acc EOF = yield (acc TL.empty) EOF
require :: Monad m => Integer -> Iteratee T.Text m ()
require n | n <= 0 = return ()
require n = continue (loop id n) where
loop acc n' (Chunks xs) = iter where
lazy = TL.fromChunks xs
len = toInteger (TL.length lazy)
iter = if len < n'
then continue (loop (acc . (TL.append lazy)) (n' len))
else yield () (toChunks (acc lazy))
loop _ _ EOF = throwError (Exc.ErrorCall "require: Unexpected EOF")
isolate :: Monad m => Integer -> Enumeratee T.Text T.Text 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 = TL.fromChunks xs
len = toInteger (TL.length lazy)
iter = if len <= n
then k (Chunks xs) >>== isolate (n len)
else let
(s1, s2) = TL.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
data Codec = Codec
{ codecName :: T.Text
, codecEncode
:: T.Text
-> (B.ByteString, Maybe (Exc.SomeException, T.Text))
, codecDecode
:: B.ByteString
-> (T.Text, Either
(Exc.SomeException, B.ByteString)
B.ByteString)
}
instance Show Codec where
showsPrec d c = showParen (d > 10) $
showString "Codec " . shows (codecName c)
encode :: Monad m => Codec
-> Enumeratee T.Text B.ByteString m b
encode codec = checkDone (continue . step) where
step k EOF = yield (Continue k) EOF
step k (Chunks xs) = loop k xs
loop k [] = continue (step k)
loop k (x:xs) = let
(bytes, extra) = codecEncode codec x
extraChunks = Chunks $ case extra of
Nothing -> xs
Just (_, text) -> text:xs
checkError k' = case extra of
Nothing -> loop k' xs
Just (exc, _) -> throwError exc
in if B.null bytes
then checkError k
else k (Chunks [bytes]) >>==
checkDoneEx extraChunks checkError
decode :: Monad m => Codec
-> Enumeratee B.ByteString T.Text m b
decode codec = checkDone (continue . step B.empty) where
step _ k EOF = yield (Continue k) EOF
step acc k (Chunks xs) = loop acc k xs
loop acc k [] = continue (step acc k)
loop acc k (x:xs) = let
(text, extra) = codecDecode codec (B.append acc x)
extraChunks = Chunks (either snd id extra : xs)
checkError k' = case extra of
Left (exc, _) -> throwError exc
Right bytes -> loop bytes k' xs
in if T.null text
then checkError k
else k (Chunks [text]) >>==
checkDoneEx extraChunks checkError
byteSplits :: B.ByteString
-> [(B.ByteString, B.ByteString)]
byteSplits bytes = loop (B.length bytes) where
loop 0 = [(B.empty, bytes)]
loop n = B.splitAt n bytes : loop (n 1)
splitSlowly :: (B.ByteString -> T.Text)
-> B.ByteString
-> (T.Text, Either
(Exc.SomeException, B.ByteString)
B.ByteString)
splitSlowly dec bytes = valid where
valid = firstValid (Prelude.map decFirst splits)
splits = byteSplits bytes
firstValid = Prelude.head . catMaybes
tryDec = tryEvaluate . dec
decFirst (a, b) = case tryDec a of
Left _ -> Nothing
Right text -> Just (text, case tryDec b of
Left exc -> Left (exc, b)
Right _ -> Right B.empty)
utf8 :: Codec
utf8 = Codec name enc dec where
name = T.pack "UTF-8"
enc text = (TE.encodeUtf8 text, Nothing)
dec bytes = case splitQuickly bytes of
Just (text, extra) -> (text, Right extra)
Nothing -> splitSlowly TE.decodeUtf8 bytes
splitQuickly bytes = loop 0 >>= maybeDecode where
required x0
| x0 .&. 0x80 == 0x00 = 1
| x0 .&. 0xE0 == 0xC0 = 2
| x0 .&. 0xF0 == 0xE0 = 3
| x0 .&. 0xF8 == 0xF0 = 4
| otherwise = 0
maxN = B.length bytes
loop n | n == maxN = Just (TE.decodeUtf8 bytes, B.empty)
loop n = let
req = required (B.index bytes n)
tooLong = first TE.decodeUtf8 (B.splitAt n bytes)
decodeMore = loop $! n + req
in if req == 0
then Nothing
else if n + req > maxN
then Just tooLong
else decodeMore
utf16_le :: Codec
utf16_le = Codec name enc dec where
name = T.pack "UTF-16-LE"
enc text = (TE.encodeUtf16LE text, Nothing)
dec bytes = case splitQuickly bytes of
Just (text, extra) -> (text, Right extra)
Nothing -> splitSlowly TE.decodeUtf16LE bytes
splitQuickly bytes = maybeDecode (loop 0) where
maxN = B.length bytes
loop n | n == maxN = decodeAll
| (n + 1) == maxN = decodeTo n
loop n = let
req = utf16Required
(B.index bytes 0)
(B.index bytes 1)
decodeMore = loop $! n + req
in if n + req > maxN
then decodeTo n
else decodeMore
decodeTo n = first TE.decodeUtf16LE (B.splitAt n bytes)
decodeAll = (TE.decodeUtf16LE bytes, B.empty)
utf16_be :: Codec
utf16_be = Codec name enc dec where
name = T.pack "UTF-16-BE"
enc text = (TE.encodeUtf16BE text, Nothing)
dec bytes = case splitQuickly bytes of
Just (text, extra) -> (text, Right extra)
Nothing -> splitSlowly TE.decodeUtf16BE bytes
splitQuickly bytes = maybeDecode (loop 0) where
maxN = B.length bytes
loop n | n == maxN = decodeAll
| (n + 1) == maxN = decodeTo n
loop n = let
req = utf16Required
(B.index bytes 1)
(B.index bytes 0)
decodeMore = loop $! n + req
in if n + req > maxN
then decodeTo n
else decodeMore
decodeTo n = first TE.decodeUtf16BE (B.splitAt n bytes)
decodeAll = (TE.decodeUtf16BE bytes, B.empty)
utf16Required :: Word8 -> Word8 -> Int
utf16Required x0 x1 = required where
required = if x >= 0xD800 && x <= 0xDBFF
then 4
else 2
x :: Word16
x = (fromIntegral x1 `shiftL` 8) .|. fromIntegral x0
utf32_le :: Codec
utf32_le = Codec name enc dec where
name = T.pack "UTF-32-LE"
enc text = (TE.encodeUtf32LE text, Nothing)
dec bs = case utf32SplitBytes TE.decodeUtf32LE bs of
Just (text, extra) -> (text, Right extra)
Nothing -> splitSlowly TE.decodeUtf32LE bs
utf32_be :: Codec
utf32_be = Codec name enc dec where
name = T.pack "UTF-32-BE"
enc text = (TE.encodeUtf32BE text, Nothing)
dec bs = case utf32SplitBytes TE.decodeUtf32BE bs of
Just (text, extra) -> (text, Right extra)
Nothing -> splitSlowly TE.decodeUtf32BE bs
utf32SplitBytes :: (B.ByteString -> T.Text)
-> B.ByteString
-> Maybe (T.Text, B.ByteString)
utf32SplitBytes dec bytes = split where
split = maybeDecode (dec toDecode, extra)
len = B.length bytes
lenExtra = mod len 4
lenToDecode = len lenExtra
(toDecode, extra) = if lenExtra == 0
then (bytes, B.empty)
else B.splitAt lenToDecode bytes
ascii :: Codec
ascii = Codec name enc dec where
name = T.pack "ASCII"
enc text = (bytes, extra) where
(safe, unsafe) = tSpanBy (\c -> ord c <= 0x7F) text
bytes = B8.pack (T.unpack safe)
extra = if T.null unsafe
then Nothing
else Just (illegalEnc name (T.head unsafe), unsafe)
dec bytes = (text, extra) where
(safe, unsafe) = B.span (<= 0x7F) bytes
text = T.pack (B8.unpack safe)
extra = if B.null unsafe
then Right B.empty
else Left (illegalDec name (B.head unsafe), unsafe)
iso8859_1 :: Codec
iso8859_1 = Codec name enc dec where
name = T.pack "ISO-8859-1"
enc text = (bytes, extra) where
(safe, unsafe) = tSpanBy (\c -> ord c <= 0xFF) text
bytes = B8.pack (T.unpack safe)
extra = if T.null unsafe
then Nothing
else Just (illegalEnc name (T.head unsafe), unsafe)
dec bytes = (T.pack (B8.unpack bytes), Right B.empty)
illegalEnc :: T.Text -> Char -> Exc.SomeException
illegalEnc name c = Exc.toException . Exc.ErrorCall $
concat [ "Codec "
, show name
, " can't encode character "
, reprChar c
]
illegalDec :: T.Text -> Word8 -> Exc.SomeException
illegalDec name w = Exc.toException . Exc.ErrorCall $
concat [ "Codec "
, show name
, " can't decode byte "
, reprWord w
]
tryEvaluate :: a -> Either Exc.SomeException a
tryEvaluate = unsafePerformIO . Exc.try . Exc.evaluate
maybeDecode:: (a, b) -> Maybe (a, b)
maybeDecode (a, b) = case tryEvaluate a of
Left _ -> Nothing
Right _ -> Just (a, b)