{-# LANGUAGE DataKinds #-}

-- | The "System.Process.Typed" module from @typed-process@, but with
-- added conduit helpers.
module Data.Conduit.Process.Typed
  ( -- * Conduit specific stuff
    createSink
  , createSinkClose
  , createSource
    -- * Running a process with logging
  , withLoggedProcess_
    -- * Reexports
  , module System.Process.Typed
  ) where

import System.Process.Typed
import qualified System.Process.Typed as P
import Data.Conduit (ConduitM, (.|), runConduit)
import qualified Data.Conduit.Binary as CB
import Control.Monad.IO.Unlift
import qualified Data.ByteString as S
import qualified Data.Conduit.List as CL
import qualified Data.ByteString.Lazy as BL
import Data.IORef (IORef, newIORef, readIORef, modifyIORef)
import Control.Exception (throwIO, catch)
import Control.Concurrent.Async (concurrently)
import System.IO (hSetBuffering, BufferMode (NoBuffering), hClose)

-- | Provide input to a process by writing to a conduit. The sink provided here
-- will leave the pipe to the child open after the stream ends. This allows the
-- sink to be used multiple times, but may result in surprising behavior. You
-- may prefer 'createSinkClose', see
-- <https://github.com/snoyberg/conduit/issues/434>.
--
-- @since 1.2.1
createSink :: MonadIO m => StreamSpec 'STInput (ConduitM S.ByteString o m ())
createSink :: StreamSpec 'STInput (ConduitM ByteString o m ())
createSink =
  (\Handle
h -> IO () -> ConduitM ByteString o m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Handle -> BufferMode -> IO ()
hSetBuffering Handle
h BufferMode
NoBuffering) ConduitM ByteString o m ()
-> ConduitM ByteString o m () -> ConduitM ByteString o m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Handle -> ConduitM ByteString o m ()
forall (m :: * -> *) o.
MonadIO m =>
Handle -> ConduitT ByteString o m ()
CB.sinkHandle Handle
h)
  (Handle -> ConduitM ByteString o m ())
-> StreamSpec 'STInput Handle
-> StreamSpec 'STInput (ConduitM ByteString o m ())
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` StreamSpec 'STInput Handle
forall (anyStreamType :: StreamType).
StreamSpec anyStreamType Handle
createPipe

-- | Like 'createSink', but closes the pipe to the child process as soon as it
-- runs out of data.
--
-- @since 1.3.5
createSinkClose :: MonadIO m => StreamSpec 'STInput (ConduitM S.ByteString o m ())
createSinkClose :: StreamSpec 'STInput (ConduitM ByteString o m ())
createSinkClose =
  (\Handle
h -> IO () -> ConduitM ByteString o m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Handle -> BufferMode -> IO ()
hSetBuffering Handle
h BufferMode
NoBuffering) ConduitM ByteString o m ()
-> ConduitM ByteString o m () -> ConduitM ByteString o m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Handle -> ConduitM ByteString o m ()
forall (m :: * -> *) o.
MonadIO m =>
Handle -> ConduitT ByteString o m ()
CB.sinkHandle Handle
h ConduitM ByteString o m ()
-> ConduitM ByteString o m () -> ConduitM ByteString o m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> IO () -> ConduitM ByteString o m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Handle -> IO ()
hClose Handle
h))
  (Handle -> ConduitM ByteString o m ())
-> StreamSpec 'STInput Handle
-> StreamSpec 'STInput (ConduitM ByteString o m ())
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` StreamSpec 'STInput Handle
forall (anyStreamType :: StreamType).
StreamSpec anyStreamType Handle
createPipe

-- | Read output from a process by read from a conduit.
--
-- @since 1.2.1
createSource :: MonadIO m => StreamSpec 'STOutput (ConduitM i S.ByteString m ())
createSource :: StreamSpec 'STOutput (ConduitM i ByteString m ())
createSource =
  (\Handle
h -> IO () -> ConduitM i ByteString m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Handle -> BufferMode -> IO ()
hSetBuffering Handle
h BufferMode
NoBuffering) ConduitM i ByteString m ()
-> ConduitM i ByteString m () -> ConduitM i ByteString m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Handle -> ConduitM i ByteString m ()
forall (m :: * -> *) i.
MonadIO m =>
Handle -> ConduitT i ByteString m ()
CB.sourceHandle Handle
h)
  (Handle -> ConduitM i ByteString m ())
-> StreamSpec 'STOutput Handle
-> StreamSpec 'STOutput (ConduitM i ByteString m ())
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` StreamSpec 'STOutput Handle
forall (anyStreamType :: StreamType).
StreamSpec anyStreamType Handle
createPipe

-- | Internal function: like 'createSource', but stick all chunks into
-- the 'IORef'.
createSourceLogged
  :: MonadIO m
  => IORef ([S.ByteString] -> [S.ByteString])
  -> StreamSpec 'STOutput (ConduitM i S.ByteString m ())
createSourceLogged :: IORef ([ByteString] -> [ByteString])
-> StreamSpec 'STOutput (ConduitM i ByteString m ())
createSourceLogged IORef ([ByteString] -> [ByteString])
ref =
    -- We do not add a cleanup action to close the handle, since in
    -- withLoggedProcess_ we attempt to read from the handle twice
    (\Handle
h ->
       (  Handle -> ConduitM i ByteString m ()
forall (m :: * -> *) i.
MonadIO m =>
Handle -> ConduitT i ByteString m ()
CB.sourceHandle Handle
h
       ConduitM i ByteString m ()
-> ConduitT ByteString ByteString m ()
-> ConduitM i ByteString m ()
forall (m :: * -> *) a b c r.
Monad m =>
ConduitT a b m () -> ConduitT b c m r -> ConduitT a c m r
.| (ByteString -> m ()) -> ConduitT ByteString ByteString m ()
forall (m :: * -> *) a. Monad m => (a -> m ()) -> ConduitT a a m ()
CL.iterM (\ByteString
bs -> IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ IORef ([ByteString] -> [ByteString])
-> (([ByteString] -> [ByteString]) -> [ByteString] -> [ByteString])
-> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef IORef ([ByteString] -> [ByteString])
ref (([ByteString] -> [ByteString])
-> ([ByteString] -> [ByteString]) -> [ByteString] -> [ByteString]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ByteString
bsByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
:))))
    )
    (Handle -> ConduitM i ByteString m ())
-> StreamSpec 'STOutput Handle
-> StreamSpec 'STOutput (ConduitM i ByteString m ())
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` StreamSpec 'STOutput Handle
forall (anyStreamType :: StreamType).
StreamSpec anyStreamType Handle
createPipe

-- | Run a process, throwing an exception on a failure exit code. This
-- will store all output from stdout and stderr in memory for better
-- error messages. Note that this will require unbounded memory usage,
-- so caveat emptor.
--
-- This will ignore any previous settings for the stdout and stderr
-- streams, and instead force them to use 'createSource'.
--
-- @since 1.2.3
withLoggedProcess_
  :: MonadUnliftIO m
  => ProcessConfig stdin stdoutIgnored stderrIgnored
  -> (Process stdin (ConduitM () S.ByteString m ()) (ConduitM () S.ByteString m ()) -> m a)
  -> m a
withLoggedProcess_ :: ProcessConfig stdin stdoutIgnored stderrIgnored
-> (Process
      stdin (ConduitM () ByteString m ()) (ConduitM () ByteString m ())
    -> m a)
-> m a
withLoggedProcess_ ProcessConfig stdin stdoutIgnored stderrIgnored
pc Process
  stdin (ConduitM () ByteString m ()) (ConduitM () ByteString m ())
-> m a
inner = (UnliftIO m -> IO a) -> m a
forall (m :: * -> *) a.
MonadUnliftIO m =>
(UnliftIO m -> IO a) -> m a
withUnliftIO ((UnliftIO m -> IO a) -> m a) -> (UnliftIO m -> IO a) -> m a
forall a b. (a -> b) -> a -> b
$ \UnliftIO m
u -> do
  IORef ([ByteString] -> [ByteString])
stdoutBuffer <- ([ByteString] -> [ByteString])
-> IO (IORef ([ByteString] -> [ByteString]))
forall a. a -> IO (IORef a)
newIORef [ByteString] -> [ByteString]
forall a. a -> a
id
  IORef ([ByteString] -> [ByteString])
stderrBuffer <- ([ByteString] -> [ByteString])
-> IO (IORef ([ByteString] -> [ByteString]))
forall a. a -> IO (IORef a)
newIORef [ByteString] -> [ByteString]
forall a. a -> a
id
  let pc' :: ProcessConfig
  stdin (ConduitM i ByteString m ()) (ConduitM i ByteString m ())
pc' = StreamSpec 'STOutput (ConduitM i ByteString m ())
-> ProcessConfig stdin stdoutIgnored (ConduitM i ByteString m ())
-> ProcessConfig
     stdin (ConduitM i ByteString m ()) (ConduitM i ByteString m ())
forall stdout stdin stdout0 stderr.
StreamSpec 'STOutput stdout
-> ProcessConfig stdin stdout0 stderr
-> ProcessConfig stdin stdout stderr
setStdout (IORef ([ByteString] -> [ByteString])
-> StreamSpec 'STOutput (ConduitM i ByteString m ())
forall (m :: * -> *) i.
MonadIO m =>
IORef ([ByteString] -> [ByteString])
-> StreamSpec 'STOutput (ConduitM i ByteString m ())
createSourceLogged IORef ([ByteString] -> [ByteString])
stdoutBuffer)
          (ProcessConfig stdin stdoutIgnored (ConduitM i ByteString m ())
 -> ProcessConfig
      stdin (ConduitM i ByteString m ()) (ConduitM i ByteString m ()))
-> ProcessConfig stdin stdoutIgnored (ConduitM i ByteString m ())
-> ProcessConfig
     stdin (ConduitM i ByteString m ()) (ConduitM i ByteString m ())
forall a b. (a -> b) -> a -> b
$ StreamSpec 'STOutput (ConduitM i ByteString m ())
-> ProcessConfig stdin stdoutIgnored stderrIgnored
-> ProcessConfig stdin stdoutIgnored (ConduitM i ByteString m ())
forall stderr stdin stdout stderr0.
StreamSpec 'STOutput stderr
-> ProcessConfig stdin stdout stderr0
-> ProcessConfig stdin stdout stderr
setStderr (IORef ([ByteString] -> [ByteString])
-> StreamSpec 'STOutput (ConduitM i ByteString m ())
forall (m :: * -> *) i.
MonadIO m =>
IORef ([ByteString] -> [ByteString])
-> StreamSpec 'STOutput (ConduitM i ByteString m ())
createSourceLogged IORef ([ByteString] -> [ByteString])
stderrBuffer) ProcessConfig stdin stdoutIgnored stderrIgnored
pc
  -- withProcessWait vs Term doesn't actually matter here, since we
  -- call checkExitCode inside regardless. But still, Wait is the
  -- safer function to use in general.
  ProcessConfig
  stdin (ConduitM () ByteString m ()) (ConduitM () ByteString m ())
-> (Process
      stdin (ConduitM () ByteString m ()) (ConduitM () ByteString m ())
    -> IO a)
-> IO a
forall (m :: * -> *) stdin stdout stderr a.
MonadUnliftIO m =>
ProcessConfig stdin stdout stderr
-> (Process stdin stdout stderr -> m a) -> m a
P.withProcessWait ProcessConfig
  stdin (ConduitM () ByteString m ()) (ConduitM () ByteString m ())
forall i i.
ProcessConfig
  stdin (ConduitM i ByteString m ()) (ConduitM i ByteString m ())
pc' ((Process
    stdin (ConduitM () ByteString m ()) (ConduitM () ByteString m ())
  -> IO a)
 -> IO a)
-> (Process
      stdin (ConduitM () ByteString m ()) (ConduitM () ByteString m ())
    -> IO a)
-> IO a
forall a b. (a -> b) -> a -> b
$ \Process
  stdin (ConduitM () ByteString m ()) (ConduitM () ByteString m ())
p -> do
    a
a <- UnliftIO m -> forall a. m a -> IO a
forall (m :: * -> *). UnliftIO m -> forall a. m a -> IO a
unliftIO UnliftIO m
u (m a -> IO a) -> m a -> IO a
forall a b. (a -> b) -> a -> b
$ Process
  stdin (ConduitM () ByteString m ()) (ConduitM () ByteString m ())
-> m a
inner Process
  stdin (ConduitM () ByteString m ()) (ConduitM () ByteString m ())
p
    let drain :: ConduitT () b m () -> IO ()
drain ConduitT () b m ()
src = UnliftIO m -> m () -> IO ()
forall (m :: * -> *). UnliftIO m -> forall a. m a -> IO a
unliftIO UnliftIO m
u (ConduitT () Void m () -> m ()
forall (m :: * -> *) r. Monad m => ConduitT () Void m r -> m r
runConduit (ConduitT () b m ()
src ConduitT () b m () -> ConduitT b Void m () -> ConduitT () Void m ()
forall (m :: * -> *) a b c r.
Monad m =>
ConduitT a b m () -> ConduitT b c m r -> ConduitT a c m r
.| ConduitT b Void m ()
forall (m :: * -> *) i o. Monad m => ConduitT i o m ()
CL.sinkNull))
    ((), ()) <- ConduitM () ByteString m () -> IO ()
forall b. ConduitT () b m () -> IO ()
drain (Process
  stdin (ConduitM () ByteString m ()) (ConduitM () ByteString m ())
-> ConduitM () ByteString m ()
forall stdin stdout stderr. Process stdin stdout stderr -> stdout
getStdout Process
  stdin (ConduitM () ByteString m ()) (ConduitM () ByteString m ())
p) IO () -> IO () -> IO ((), ())
forall a b. IO a -> IO b -> IO (a, b)
`concurrently`
                ConduitM () ByteString m () -> IO ()
forall b. ConduitT () b m () -> IO ()
drain (Process
  stdin (ConduitM () ByteString m ()) (ConduitM () ByteString m ())
-> ConduitM () ByteString m ()
forall stdin stdout stderr. Process stdin stdout stderr -> stderr
getStderr Process
  stdin (ConduitM () ByteString m ()) (ConduitM () ByteString m ())
p)
    Process
  stdin (ConduitM () ByteString m ()) (ConduitM () ByteString m ())
-> IO ()
forall (m :: * -> *) stdin stdout stderr.
MonadIO m =>
Process stdin stdout stderr -> m ()
checkExitCode Process
  stdin (ConduitM () ByteString m ()) (ConduitM () ByteString m ())
p IO () -> (ExitCodeException -> IO ()) -> IO ()
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` \ExitCodeException
ece -> do
      [ByteString] -> [ByteString]
stdout <- IORef ([ByteString] -> [ByteString])
-> IO ([ByteString] -> [ByteString])
forall a. IORef a -> IO a
readIORef IORef ([ByteString] -> [ByteString])
stdoutBuffer
      [ByteString] -> [ByteString]
stderr <- IORef ([ByteString] -> [ByteString])
-> IO ([ByteString] -> [ByteString])
forall a. IORef a -> IO a
readIORef IORef ([ByteString] -> [ByteString])
stderrBuffer
      ExitCodeException -> IO ()
forall e a. Exception e => e -> IO a
throwIO ExitCodeException
ece
        { eceStdout :: ByteString
eceStdout = [ByteString] -> ByteString
BL.fromChunks ([ByteString] -> ByteString) -> [ByteString] -> ByteString
forall a b. (a -> b) -> a -> b
$ [ByteString] -> [ByteString]
stdout []
        , eceStderr :: ByteString
eceStderr = [ByteString] -> ByteString
BL.fromChunks ([ByteString] -> ByteString) -> [ByteString] -> ByteString
forall a b. (a -> b) -> a -> b
$ [ByteString] -> [ByteString]
stderr []
        }
    a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return a
a