Process.hs

{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE MultiParamTypeClasses #-}

-- | Reading from the process.

module Data.Conduit.Shell.Process
  (-- * Running scripts
   run
   -- * Conduit types
  ,text
  ,bytes
  -- * General conduits
  ,conduit
  ,conduitEither
  -- * Running processes
  ,Data.Conduit.Shell.Process.shell
  ,Data.Conduit.Shell.Process.proc
  ,($|)
  ,Segment
  ,ProcessException(..)
  ,ToChunk(..)
  ,tryS
  )
  where

import           Control.Applicative
import           Control.Concurrent.Async
import           Control.Exception
import           Control.Monad
import           Control.Monad.IO.Class
import           Data.ByteString (ByteString)
import qualified Data.ByteString as S
import           Data.Conduit
import           Data.Conduit.Binary
import qualified Data.Conduit.List as CL
import           Conduit (MonadThrow)
import           Data.Conduit.Text (encodeUtf8, decodeUtf8)
import           Data.Text (Text)
import           Data.Typeable
import           System.Exit
import           System.IO
import           System.Posix.IO (createPipe, fdToHandle)
import           System.Process hiding (createPipe)
import           UnliftIO (MonadUnliftIO, unliftIO, askUnliftIO)

-- | A pipeable segment. Either a conduit or a process.
data Segment m r
  = SegmentConduit (ConduitM ByteString (Either ByteString ByteString) m r)
  | SegmentProcess (Handles -> m r)

instance MonadIO m => Monad (Segment m) where
  return = SegmentConduit . return
  SegmentConduit c >>= f =
    SegmentProcess (conduitToProcess c) >>=
    f
  SegmentProcess f >>= g =
    SegmentProcess
      (\handles ->
         do x <- f handles
            case g x of
              SegmentConduit c ->
                conduitToProcess c handles
              SegmentProcess p -> p handles)

instance MonadIO m => Functor (Segment m) where
  fmap = liftM

instance MonadIO m => Applicative (Segment m) where
  (<*>) = ap; pure = return

instance MonadUnliftIO m => Alternative (Segment m) where
  this <|> that =
    do ex <- tryS this
       case ex of
         Right x -> pure x
         Left (_ :: ProcessException) -> that
  empty = throw ProcessEmpty

-- | Try something in a segment.
tryS :: (Exception e, MonadUnliftIO m) => Segment m r -> Segment m (Either e r)
tryS s =
  case s of
    SegmentConduit c -> SegmentConduit (tryC c)
    SegmentProcess f -> SegmentProcess $ (\h -> do
                                             u <- askUnliftIO
                                             liftIO $ try $ unliftIO u (f h))

instance MonadIO m => MonadIO (Segment m) where
  liftIO x = SegmentProcess (const $ liftIO  x)

-- | Process handles: @stdin@, @stdout@, @stderr@
data Handles =
  Handles Handle
          Handle
          Handle

-- | Process running exception.
data ProcessException
  = ProcessException CreateProcess
                     ExitCode
  | ProcessEmpty
  deriving (Typeable)

instance Exception ProcessException

instance Show ProcessException where
  show ProcessEmpty = "empty process"
  show (ProcessException cp ec) =
    concat
      [ "The "
      , case cmdspec cp of
          ShellCommand s -> "shell command " ++ show s
          RawCommand f args -> "raw command: " ++ unwords (f : map show args)
      , " returned a failure exit code: "
      , case ec of
          ExitFailure i -> show i
          _ -> show ec
      ]

-- | Convert a process or a conduit to a segment.
class ToSegment m a  where
  type SegmentResult m a
  toSegment :: a -> Segment m (SegmentResult m a)

instance ToSegment m (Segment m r) where
  type SegmentResult m (Segment m r) = r
  toSegment = id

instance (a ~ ByteString, ToChunk b, Monad m) =>
         ToSegment m (ConduitT a b m r) where
  type SegmentResult m (ConduitT a b m r) = r
  toSegment f = SegmentConduit (f `fuseUpstream` CL.map toChunk)

instance MonadIO m => ToSegment m CreateProcess where
  type SegmentResult m CreateProcess = ()
  toSegment = liftProcess

-- | Used to allow outputting stdout or stderr.
class ToChunk a  where
  toChunk :: a -> Either ByteString ByteString

instance ToChunk ByteString where
  toChunk = Left

instance ToChunk (Either ByteString ByteString) where
  toChunk = id

-- | Run a shell command.
shell :: MonadIO m => String -> Segment m ()
shell = liftProcess . System.Process.shell

-- | Run a process command.
proc :: MonadIO m => String -> [String] -> Segment m ()
proc name args = liftProcess (System.Process.proc name args)

-- | Run a segment.
run :: MonadIO m => Segment m r -> m r
run (SegmentConduit c) = run (SegmentProcess (conduitToProcess c))
run (SegmentProcess p) = p (Handles stdin stdout stderr)

-- | Fuse two segments (either processes or conduits).
($|) :: MonadUnliftIO m => Segment m () -> Segment m b -> Segment m b
x $| y = x `fuseSegment` y

infixl 0 $|

-- | Work on the stream as 'Text' values from UTF-8.
text
  :: (r ~ (), MonadThrow m)
  => ConduitT Text Text m r -> Segment m r
text conduit' = bytes (decodeUtf8 .| conduit' .| encodeUtf8)

-- | Lift a conduit into a segment.
bytes
  :: (a ~ ByteString, Monad m)
  => ConduitT a ByteString m r -> Segment m r
bytes f = SegmentConduit (f `fuseUpstream` CL.map toChunk)

-- | Lift a conduit into a segment.
conduit
  :: (a ~ ByteString, Monad m)
  => ConduitT a ByteString m r -> Segment m r
conduit f = SegmentConduit (f `fuseUpstream` CL.map toChunk)

-- | Lift a conduit into a segment, which can yield stderr.
conduitEither
  :: (a ~ ByteString, Monad m)
  => ConduitT a (Either ByteString ByteString) m r -> Segment m r
conduitEither f = SegmentConduit (f `fuseUpstream` CL.map toChunk)

-- | Lift a process into a segment.
liftProcess :: MonadIO m => CreateProcess -> Segment m ()
liftProcess cp =
  SegmentProcess
    (\(Handles inh outh errh) ->
        let config =
              cp
              { std_in = UseHandle inh
              , std_out = UseHandle outh
              , std_err = UseHandle errh
              , close_fds = True
              }
        in
          liftIO $ do
              (Nothing, Nothing, Nothing, ph) <- createProcess_ "liftProcess" config
              ec <- waitForProcess ph
              case ec of
                ExitSuccess -> return ()
                _ -> throwIO (ProcessException cp ec))

-- | Convert a conduit to a process.
conduitToProcess :: MonadIO m => ConduitT ByteString (Either ByteString ByteString) m r
                 -> (Handles -> m r)
conduitToProcess c (Handles inh outh errh) =
  runConduit $ sourceHandle inh .| c `fuseUpstream` sinkHandles outh errh

-- | Sink everything into the two handles.
sinkHandles ::
            MonadIO m
            => Handle
            -> Handle
            -> ConduitT (Either ByteString ByteString) Void m ()
sinkHandles out err =
  CL.mapM_
    (\ebs ->
        liftIO $ case ebs of
          Left bs -> S.hPut out bs
          Right bs -> S.hPut err bs)

-- | Create a pipe.
createHandles :: IO (Handle, Handle)
createHandles =
  mask_
    (do (inFD, outFD) <- createPipe
        x <- fdToHandle inFD
        y <- fdToHandle outFD
        hSetBuffering x NoBuffering
        hSetBuffering y NoBuffering
        return (x, y))

-- | Fuse two processes.
fuseProcess :: MonadUnliftIO m => (Handles -> m ()) -> (Handles -> m r) -> (Handles -> m r)
fuseProcess left right (Handles in1 out2 err) = do
  u <- askUnliftIO
  (in2, out1) <- liftIO createHandles
  liftIO $ runConcurrently
    (Concurrently ((unliftIO u $ left (Handles in1 out1 err)) `finally` hClose out1) *>
     Concurrently ((unliftIO u $ right (Handles in2 out2 err)) `finally` hClose in2))

-- | Fuse two conduits.
fuseConduit
  :: Monad m
  => ConduitT ByteString (Either ByteString ByteString) m ()
  -> ConduitT ByteString (Either ByteString ByteString) m r
  -> ConduitT ByteString (Either ByteString ByteString) m r
fuseConduit left right = left .| getZipConduit right'
  where
    right' =
      ZipConduit (CL.filter isRight) *>
      ZipConduit (CL.mapMaybe (either (const Nothing) Just) .| right)
    isRight Right {} = True
    isRight Left {} = False

-- | Fuse a conduit with a process.
fuseConduitProcess
  :: MonadUnliftIO m
  => ConduitT ByteString (Either ByteString ByteString) m ()
  -> (Handles -> m r)
  -> (Handles -> m r)
fuseConduitProcess left right (Handles in1 out2 err) = do
  u <- askUnliftIO
  (in2, out1) <- liftIO createHandles
  liftIO $ runConcurrently
    (Concurrently
       ((unliftIO u $ runConduit $ sourceHandle in1 .| left .| sinkHandles out1 err) `finally`
        hClose out1) *>
     Concurrently ((unliftIO u $ right (Handles in2 out2 err)) `finally` hClose in2))

-- | Fuse a process with a conduit.
fuseProcessConduit
  :: MonadUnliftIO m
  => (Handles -> m ())
  -> ConduitT ByteString (Either ByteString ByteString) m r
  -> (Handles -> m r)
fuseProcessConduit left right (Handles in1 out2 err) = do
  u <- askUnliftIO
  (in2, out1) <- liftIO createHandles
  liftIO $ runConcurrently
    (Concurrently ((unliftIO u $ left (Handles in1 out1 err)) `finally` hClose out1) *>
     Concurrently
       ((unliftIO u $ runConduit $
         sourceHandle in2 .| right `fuseUpstream` sinkHandles out2 err) `finally`
        hClose in2))

-- | Fuse one segment with another.
fuseSegment :: MonadUnliftIO m => Segment m () -> Segment m r -> Segment m r
SegmentConduit x `fuseSegment` SegmentConduit y =
  SegmentConduit (fuseConduit x y)
SegmentConduit x `fuseSegment` SegmentProcess y =
  SegmentProcess (fuseConduitProcess x y)
SegmentProcess x `fuseSegment` SegmentConduit y =
  SegmentProcess (fuseProcessConduit x y)
SegmentProcess x `fuseSegment` SegmentProcess y =
  SegmentProcess (fuseProcess x y)