summaryrefslogtreecommitdiff
path: root/lib/ProcessUtils.hs
blob: 492d666fb80c17dc62f1a803bdaf5043d4beecfe (plain)
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
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
{-# LANGUAGE CPP #-}
{-# LANGUAGE OverloadedStrings #-}
module ProcessUtils
    ( ExitCode(ExitFailure,ExitSuccess)
    , systemEnv
    , readPipe
    , readProcessWithErrorH
    ) where

import GHC.IO.Exception ( ioException, IOErrorType(..) )
import System.FilePath
import System.Directory
import System.Process
import System.Posix.Signals
import System.Process.Internals (runGenProcess_,defaultSignal)
import System.Environment
import Data.Maybe ( isNothing, fromMaybe )
import System.IO.Error ( mkIOError, ioeSetErrorString ) 
import System.Exit ( ExitCode(..) )
import System.IO
import Control.Applicative
import Control.Exception (bracket)
import qualified Data.ByteString as S


-- | systemEnv
-- This is like System.Process.system except that it lets you set
-- some environment variables.
systemEnv :: [(String, String)] -> String -> IO ExitCode
systemEnv _    ""  = 
    ioException (ioeSetErrorString (mkIOError InvalidArgument "system" Nothing Nothing) "null command")
systemEnv vars cmd = do
    env0 <- getEnvironment
    let env1 = filter (isNothing . flip lookup vars . fst) env0
        env = vars ++ env1
    syncProcess "system" $ (shell cmd) {env=Just env}
 where
    -- This is a non-exported function from System.Process
    syncProcess fun c = do
      -- The POSIX version of system needs to do some manipulation of signal
      -- handlers.  Since we're going to be synchronously waiting for the child,
      -- we want to ignore ^C in the parent, but handle it the default way
      -- in the child (using SIG_DFL isn't really correct, it should be the
      -- original signal handler, but the GHC RTS will have already set up
      -- its own handler and we don't want to use that).
      old_int  <- installHandler sigINT  Ignore Nothing
      old_quit <- installHandler sigQUIT Ignore Nothing
      (_,_,_,p) <- runGenProcess_ fun c
            (Just defaultSignal) (Just defaultSignal)
      r <- waitForProcess p
      _ <- installHandler sigINT  old_int Nothing
      _ <- installHandler sigQUIT old_quit Nothing
      return r

-- | readPipe
--
-- This is like System.Process.readProcess but instead of a single
-- command with arguments, you pass it a list which it then executes
-- piping the the output from the prior command into the input for
-- the next command. Like readProcess the next parameter is a string
-- which is sent as input into the initial standard in.
--
--  > readPipe [("grep",["HostConfig"]),("wc",["-l"]) =<< readFile "input"
--
-- is equivalent to shell code:
--
--  > grep HostConfig input | wc -l
--
readPipe :: [(FilePath,[String])] -> String -> IO String
readPipe [(cmd,args)] input = readProcess cmd args input
readPipe ((cmd,args):xs) stdin0 = do
    let p = (shell cmd) { std_in = CreatePipe
                        , std_out = CreatePipe
                        , std_err = CreatePipe
                        , cmdspec = RawCommand cmd args
                        }
    (Just sinh,Just sout,Just serr, ph) <- createProcess p
    hPutStr sinh stdin0
    hClose sinh
    readPipe0 xs sout
    where
        readPipe0 :: [(FilePath,[String])] -> Handle -> IO String
        readPipe0 ((cmd,args):xs) stdin0 = do
            let p = (shell cmd)
                        { std_in = UseHandle stdin0
                        , std_out = CreatePipe
                        , std_err = CreatePipe
                        , cmdspec = RawCommand cmd args
                        }
            (Nothing,Just sout,Just serr, ph) <- createProcess p
            readPipe0 xs sout
        readPipe0 [] h = hGetContents h


readProcessWithErrorH :: FilePath -> [String] -> String -> Handle -> IO String
readProcessWithErrorH cmd args stdin erH = do
    let p = (shell cmd) { std_in = CreatePipe
                        , std_out = CreatePipe
                        , std_err = UseHandle erH
                        , cmdspec = RawCommand cmd args
                        }
    bracket (createProcess p) 
            (\(Just sinh,Just sout,_, ph) ->
                mapM_ hClose [sinh,sout]         )
            (\(Just sinh,Just sout,_, ph) -> do
                hPutStr sinh stdin
                hClose sinh
                hGetContents sout                )

-- | run an external command.
--
-- Arguments:
--
--   [ cmd ]   command to run
--
--   [ input ] input to command (if any)
--
-- Returns:
--
--   [ code ]   exit code returned by command
--
--   [ output ] captured stdout of command
--
--   [ err ]    captured stderr of command
runExternal :: String -> Maybe String -> IO (Int,(S.ByteString,S.ByteString))
runExternal cmd input = do
    cwd <- getCurrentDirectory
    putStr $ takeFileName cwd ++ "> " ++ cmd
                    ++ case input of
                            Nothing -> "\n"
                            Just s  -> " <<EOF\n" ++ s ++ "EOF\n"
    let p = (shell cmd) { std_in  = maybe Inherit (const CreatePipe) input
                        , std_out = CreatePipe
                        , std_err = CreatePipe }
    (mb_in,mb_out,mb_err,process) <- createProcess p
    case mb_in of
        Nothing -> return ()
        Just hin -> hPutStr hin (fromMaybe "" input) >> hClose hin
    out <- case mb_out of
            Just hout -> S.hGetContents hout
            Nothing   -> return ""
    err <- case mb_err of
            Just herr -> S.hGetContents herr
            Nothing   -> return ""
    code <- do
        c <- waitForProcess process
        return $ case c of
                    ExitSuccess   -> 0
                    ExitFailure x -> x
    return (code,(out,err))