ackぽい再帰的grepを実装してみました

ソースコード検索の ack がよさげな件: blog.bulknews.netなどで紹介されているackぽいコマンドをHaskellで実装してみました。関数を繋いで繋いで、できるだけHaskellぽい実装になるように気をつけてみました。特にmain関数とか。
ファイル: hck.hs

module Main (main) where

import Control.Monad (when)
import Data.List (isSuffixOf)
import System (getArgs)
import System.Exit (exitWith, ExitCode(ExitFailure))
import System.Directory (getDirectoryContents, doesDirectoryExist)
import System.IO (readFile)
import Text.Printf (printf)
import Text.Regex (mkRegex, matchRegex, subRegex, Regex)

getDirectoryFilesRecursive :: FilePath -> IO [FilePath]
getDirectoryFilesRecursive dir =
  return . concat =<< mapM trav . notDot =<< getDirectoryContents dir
  where
    notDot :: [FilePath] -> [FilePath]
    notDot = filter (\ f -> all (/= f) [".", ".."])

    trav :: FilePath -> IO [FilePath]
    trav file = do
      let path = dir ++ "/" ++ file
      isDir <- doesDirectoryExist path
      if isDir then getDirectoryFilesRecursive path
               else return [path]

ignoreExt :: [FilePath] -> [FilePath]
ignoreExt = 
  let exts = [".o", ".exe", ".hi", "~", ".bak", "/CVS", "/RCS", "/.svn"]
  in filter (\ f -> not $ any (flip isSuffixOf f) exts)

type FileContents = (FilePath, [(Int, String)])

readContentsFromFiles :: [FilePath] -> IO [FileContents]
readContentsFromFiles files =
  return . zip files . map (zip [1 ..] . lines) =<< mapM readFile files

grep :: Regex -> [FileContents] -> [FileContents]
grep pat = 
  filter (\ (f, ls) -> not $ null ls) .
    map (\ (f, ls) -> (f, filter match ls))
  where
    match :: (Int, String) -> Bool
    match (n, l) =
      case matchRegex pat l of
        Just s  -> True
        Nothing -> False

format :: Regex -> [FileContents] -> [String]
format pat =
  concatMap (\ (f, ls) -> green f : map format' ls)
  where
    format' :: (Int, String) -> String
    format' (n, l) = printf "%6d:%s" n $ yellow l

    green :: String -> String
    green s = "\x1b[1;32;40m" ++ s ++ "\x1b[0m"

    yellow :: String -> String
    yellow l = subRegex pat l "\x1b[30;43m\\1\x1b[0m"

main :: IO ()
main = do
  (pat, dir) <- parseArgs =<< getArgs
  mapM_ putStrLn . format pat . grep pat =<<
    readContentsFromFiles . ignoreExt =<<
    getDirectoryFilesRecursive dir
  where
    usage :: IO ()
    usage = do
      putStrLn "usage: hck [pattern] [dir]"
      exitWith (ExitFailure 1)

    parseArgs :: [String] -> IO (Regex, FilePath)
    parseArgs args = do
      when (null args) $ usage
      let pat = mkRegex $ "(" ++ args !! 0 ++ ")"
          dir = if (length args > 1) then args !! 1 else "."
      return (pat, dir)

-- Local Variables:
-- compile-command: "ghc -W -fno-warn-unused-matches --make -o hck hck"
-- End:

BUGS: オプションがまったく実装されてません。
参照: WIN32OLEでiPod Shuffleもどき - 趣味的にっき