/
Utils.hs
79 lines (66 loc) · 2.5 KB
/
Utils.hs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
module Utils
( strict
, exec_action
, modifyIORef'
, counter
, progress
, mayunzip
, reader
, writer
, readerC
, writerC
) where
import qualified Data.ByteString as S
import qualified Data.ByteString.Lazy as L
import qualified Data.Conduit.Binary as CB -- bytes
import Data.Conduit
import Data.Conduit.Internal
import Data.IORef
import Data.List (isSuffixOf)
import Control.Monad
import Control.Monad.Trans
import Data.Conduit.Zlib (ungzip)
import System.IO
strict :: L.ByteString -> S.ByteString
strict = S.concat . L.toChunks
{-| exec_action executes an action for every element seen and passes it on -}
exec_action :: (Monad m) => (a -> m ()) -> Conduit a m a
exec_action act = await >>= maybe (return ()) (\s -> PipeM (act s >> return (yield s >> exec_action act)))
{-| strict version of modifyIORef -}
modifyIORef' :: IORef a -> (a -> a) -> IO ()
modifyIORef' ref f = do
x <- readIORef ref
let x' = f x
x' `seq` writeIORef ref x'
{-| A counter will increment the IORef ref every time it sees an input -}
counter :: (Monad (t IO), Num a1, MonadTrans t) => IORef a1 -> Conduit a (t IO) a
counter ref = exec_action . const . lift $ modifyIORef' ref (+1)
progress :: (Integral a, MonadIO (t IO), MonadTrans t) => a -> (t1 -> Integer) -> IO (t1 -> t IO ())
progress totalSize sizef = do
partialref <- newIORef (0 :: Integer)
return $ \s -> do
v0 <- lift (readIORef partialref)
lift $ modifyIORef' partialref (+ sizef s)
v1 <- lift (readIORef partialref)
when (roundP v0 /= roundP v1)
(liftIO $ putStrLn $ concat ["Finished ", show . roundP $ v1, "%"])
where
roundP :: Integer -> Integer
roundP v = 5 * (round $ fromIntegral v / fromIntegral totalSize * (20.0 :: Double))
transformif :: Monad m => Bool -> Pipe l a a r m r -> Pipe l a a r m r
transformif cond trans
| cond = trans
| otherwise = idP
mayunzip :: (Monad m, MonadUnsafeIO m, MonadThrow m) => String -> Conduit S.ByteString m S.ByteString
mayunzip finput = transformif ("gz" `isSuffixOf` finput) ungzip
reader :: String -> IO L.ByteString
reader "-" = L.getContents
reader inputf = L.readFile inputf
writer :: String -> (L.ByteString -> IO ())
writer "-" = L.putStr
writer outputf = L.writeFile outputf
readerC :: (MonadIO m, MonadResource m) => String -> GSource m S.ByteString
readerC "-" = CB.sourceHandle stdin
readerC inputf = CB.sourceFile inputf
writerC "-" = CB.sinkHandle stdout
writerC outputf = CB.sinkFile outputf