darcs-benchmark-0.1.6: Shellish.hs

module Shellish( cd, pwd, run, (#), withCurrentDirectory, strictGetContents,
                 shellish, silently, verbosely, Command, sub, echo, echo_n,
                 which, canonize, resetMemoryUsed, recordMemoryUsed,
                 memoryUsed, rm_rf, rm_f, whenM, test_e, test_f, test_d,
                 resetTimeUsed, timeUsed, mv, ls, echo_err, echo_n_err )
    where

import Data.List
import System.IO
import System.Directory
import System.Exit
import Control.Monad.State
import Control.Monad.Error
import Data.Time.Clock( getCurrentTime, diffUTCTime, UTCTime(..) )
import Control.Exception ( bracket, evaluate )
import Control.Monad ( when )
import qualified Data.ByteString.Char8 as B
import System.Process( runInteractiveProcess, waitForProcess )
import qualified System.IO.UTF8 as UTF8


-- TODO maxMem is sort of a layering violation here... but who cares.
data St = St { sCode :: Int, sStderr :: B.ByteString , sStdout :: B.ByteString,
               sVerbose :: Bool, maxMem :: Int, startTime :: UTCTime }

type Command a = StateT St IO a

cd :: String -> Command ()
cd = liftIO . setCurrentDirectory

mv :: String -> String -> Command ()
mv a b = liftIO $ renameFile a b

ls :: String -> Command [String]
ls dir = liftIO $ (\\ [".", ".."]) `fmap` getDirectoryContents dir

pwd :: Command String
pwd = liftIO $ getCurrentDirectory

echo, echo_n, echo_err, echo_n_err :: String -> Command ()
echo = liftIO . UTF8.putStrLn
echo_n = liftIO . (>> hFlush System.IO.stdout) . UTF8.putStr
echo_err = liftIO . UTF8.hPutStrLn stderr
echo_n_err = liftIO . (>> hFlush stderr) . UTF8.hPutStr stderr

which :: String -> Command (Maybe String)
which = liftIO . findExecutable

canonize :: String -> Command String
canonize = liftIO . canonicalizePath

whenM :: Monad m => m Bool -> m () -> m ()
whenM c a = do res <- c
               when res a

test_e :: String -> Command Bool
test_e f = liftIO $ do
             dir <- doesDirectoryExist f
             file <- doesFileExist f
             return $ file || dir

test_f :: String -> Command Bool
test_f = liftIO . doesFileExist

test_d :: String -> Command Bool
test_d = liftIO . doesDirectoryExist

rm_rf :: String -> Command ()
rm_rf f = whenM (test_e f) $ liftIO $ removeDirectoryRecursive f

rm_f :: String -> Command ()
rm_f f = whenM (test_e f) $ liftIO $ removeFile f

run :: String -> [String] -> Command String
run cmd args = do
    (_,outH,errH,procH) <- liftIO $ runInteractiveProcess cmd args Nothing Nothing
    st <- get
    res <- liftIO $ B.hGetContents outH
    errs <- liftIO $ B.hGetContents errH
    ex <- liftIO $ waitForProcess procH
    when (sVerbose st) $ do
                 liftIO $ B.putStr res
                 liftIO $ B.putStr errs
    case ex of
      ExitSuccess -> return ()
      ExitFailure n -> fail $ "command " ++ cmd ++ " " ++ show args
                         ++ " failed with exit code " ++ show n
    put $ st { sCode = 0, sStderr = errs, sStdout = res }
    return $ B.unpack res

(#) :: String -> [String] -> Command String
cmd # args = run cmd args

silently :: Command a -> Command a
silently a = do
  x <- get
  put $ x { sVerbose = False }
  r <- a
  put x
  return r

verbosely :: Command a -> Command a
verbosely a = do
  x <- get
  put $ x { sVerbose = True }
  r <- a
  put x
  return r

sub :: Command a -> Command a
sub a = do
  -- TODO save environment as well?
  dir <- liftIO $ getCurrentDirectory
  r <- a `catchError` (\e -> (liftIO $ setCurrentDirectory dir) >> throwError e)
  liftIO $ setCurrentDirectory dir
  return r

shellish :: MonadIO m => Command a -> m a
shellish a = do
  dir <- liftIO $ getCurrentDirectory
  time <- liftIO getCurrentTime
  r <- liftIO $ evalStateT a $ empty time
  liftIO $ setCurrentDirectory dir
  return r
      where empty t = St { sCode = 0, sStderr = B.empty,
                           sStdout = B.empty, sVerbose = True,
                           maxMem = 0, startTime = t }

strictGetContents :: Handle -> IO String
strictGetContents h =
    do res <- hGetContents h
       evaluate (length res)
       return res

withCurrentDirectory :: FilePath -> IO a -> IO a
withCurrentDirectory name m =
    bracket
        (do cwd <- getCurrentDirectory
            when (name /= "") (setCurrentDirectory name)
            return cwd)
        (\oldwd -> setCurrentDirectory oldwd `catch` (\_ -> return ()))
        (const m)

recordMemoryUsed :: Int -> Command ()
recordMemoryUsed u = do
  x <- get
  put x { maxMem = maximum [maxMem x, u] }

resetMemoryUsed :: Command ()
resetMemoryUsed = do
  x <- get
  put x { maxMem = 0 }

resetTimeUsed :: Command ()
resetTimeUsed = do
  x <- get
  t <- liftIO getCurrentTime
  put x { startTime = t }

memoryUsed :: Command Int
memoryUsed = do x <- get
                return $ maxMem x

timeUsed :: Command Float
timeUsed = do x <- get
              t <- liftIO getCurrentTime
              return $ realToFrac $ diffUTCTime t (startTime x)