Iteratee で progress_display を書いてみた

Haskell

Iteratee を始めようと思いこちらを読んでいたのですが，
Lazy I/O must go! - Iteratee: 列挙ベースのI/O - 純粋関数型雑記帳
気になるフレーズが．

streamToFile "hoge" `enumPair` throbber とすることにより、ファイル書き込みに簡単に進捗表示をつけることができるようになります。その他、時間のかかるような処理にプログレスバー表示を取り付けたりするのも簡単です。

http://d.hatena.ne.jp/tanakh/20100824#p1

進捗表示？プログレスバー表示？
…
…
…
それは progress_display ではありませんか？
※ progress_display について

書いてみた

Iteratee like なパッケージとしては，見た目に分かりやすかった enumerator を使いました．
Progress.hs

module Boost.Progress (progressDisplay) where

progressDisplay :: MonadIO m
  => Integer -- expected count
  -> Handle  -- output handle
  -> String  -- leading strings
  -> String
  -> String
  -> Iteratee a m ()

progressDisplay に適当なパラメータを与え，他の Iteratee に zip して使います．
Main.hs

import Control.Exception
import Control.Monad.IO.Class
import qualified Data.Enumerator as E
import qualified Data.Enumerator.List as EL
import System.IO
import Boost.Progress

streamToHandle :: MonadIO m => Handle -> E.Iteratee Char m ()
streamToHandle h = E.continue go
  where
    go E.EOF = E.yield () E.EOF
    go (E.Chunks xs) = do
      liftIO $ hPutStr h xs
      E.continue go

main :: IO ((), ())
main = bracket (openFile "hoge" WriteMode) hClose $ \h ->
  E.run_ $
  E.enumList 64 (replicate 1000000 '\NUL') E.$$
  streamToHandle h `EL.zip` progressDisplay 1000000 stdout "\n" "" ""

出力

0%   10   20   30   40   50   60   70   80   90   100%
|----|----|----|----|----|----|----|----|----|----|
***************************************************

実装

boost::progress_display の実装に合わせています．
Progress.hs

module Boost.Progress (progressDisplay) where

import Control.Monad
import Control.Monad.IO.Class
import qualified Data.Enumerator as E
import System.IO

progressDisplay :: MonadIO m
  => Integer -- expected count
  -> Handle  -- output handle
  -> String  -- leading strings
  -> String
  -> String
  -> E.Iteratee a m ()
progressDisplay n h s1 s2 s3 = start >> E.continue (go 0 0)
  where
    expcnt = if n <= 0 then 1 else n
    start = liftIO $ do
      hPutStrLn h $
        s1 ++ "0%   10   20   30   40   50   60   70   80   90   100%\n" ++
        s2 ++ "|----|----|----|----|----|----|----|----|----|----|"
      hFlush h
      hPutStr h s3
    go _ _ E.EOF = E.yield () E.EOF
    go cnt tic (E.Chunks xs) = do
      let
        cnt' = cnt + fromIntegral (length xs)
        tic' = floor (fromIntegral cnt' / fromIntegral expcnt * 50)
      liftIO $ do
        replicateM_ (tic' - tic) $ do
          hPutChar h '*'
          hFlush h
        when (cnt' == expcnt) $ do
          hPutStrLn h "*"
          hFlush h
      if (cnt' < expcnt) then
        E.continue (go cnt' tic')
      else
        E.yield () E.EOF

まとめ

遅くなりましたが，progress_display の追悼記事ということにします．
簡単に並置合成できるという Iteratee の利点も少し分かったような気がしますね．