{-# LANGUAGE BangPatterns #-}

{- | This module defines functions and type classes
 for making the syntax more succint.
-}
module Procex.Quick (
  (<!|),
  (<<<),
  (<|),
  (|!>),
  (|>),
  capture,
  captureNoThrow,
  captureLazy,
  captureLazyNoThrow,
  captureErr,
  captureErrNoThrow,
  captureErrLazy,
  captureErrLazyNoThrow,
  captureFd,
  captureFdNoThrow,
  captureFdLazy,
  captureFdLazyNoThrow,
  pipeArgStrIn,
  mq,
  QuickCmd (..),
  QuickCmdArg (..),
  ToByteString (..),
) where

import Control.Concurrent.Async (Async)
import Control.DeepSeq (force)
import qualified Data.ByteString as BS
import Data.ByteString.Lazy (ByteString)
import qualified Data.ByteString.Lazy as B
import qualified Data.ByteString.Lazy.UTF8 as B
import Data.Foldable (foldl')
import Procex.Core
import Procex.Process
import System.IO (Handle, hClose)
import System.IO.Unsafe (unsafeInterleaveIO)
import System.Posix.IO (handleToFd)
import System.Posix.Process (ProcessStatus)
import System.Posix.Types (Fd)

-- | A helper class to convert to bytestrings with UTF-8 encoding
class ToByteString a where
  toByteString :: a -> B.ByteString

instance a ~ Char => ToByteString [a] where
  toByteString :: [a] -> ByteString
toByteString = String -> ByteString
B.fromString

instance ToByteString B.ByteString where
  toByteString :: ByteString -> ByteString
toByteString = forall a. a -> a
id

instance ToByteString BS.ByteString where
  toByteString :: ByteString -> ByteString
toByteString = ByteString -> ByteString
B.fromStrict

-- | If a type implements this, you can pass it to 'mq'.
class QuickCmdArg a where
  quickCmdArg :: a -> Cmd -> Cmd

-- | A helper class to allow lightweight syntax for executing commands
class QuickCmd a where
  quickCmd :: Cmd -> a

-- | UTF-8 encoded
instance QuickCmdArg String where
  quickCmdArg :: String -> Cmd -> Cmd
quickCmdArg String
s = ByteString -> Cmd -> Cmd
passArg forall a b. (a -> b) -> a -> b
$ String -> ByteString
B.fromString String
s

instance QuickCmdArg [String] where
  quickCmdArg :: [String] -> Cmd -> Cmd
quickCmdArg = (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c. (a -> b -> c) -> b -> a -> c
flip) forall a. QuickCmdArg a => a -> Cmd -> Cmd
quickCmdArg

instance QuickCmdArg [ByteString] where
  quickCmdArg :: [ByteString] -> Cmd -> Cmd
quickCmdArg = (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c. (a -> b -> c) -> b -> a -> c
flip) forall a. QuickCmdArg a => a -> Cmd -> Cmd
quickCmdArg

instance QuickCmdArg [(Cmd -> Cmd)] where
  quickCmdArg :: [Cmd -> Cmd] -> Cmd -> Cmd
quickCmdArg = (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c. (a -> b -> c) -> b -> a -> c
flip) forall a. QuickCmdArg a => a -> Cmd -> Cmd
quickCmdArg

instance (a ~ Fd, b ~ Cmd) => QuickCmdArg [(a, (Handle -> IO b) -> IO b)] where
  quickCmdArg :: [(a, (Handle -> IO b) -> IO b)] -> Cmd -> Cmd
quickCmdArg = (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c. (a -> b -> c) -> b -> a -> c
flip) forall a. QuickCmdArg a => a -> Cmd -> Cmd
quickCmdArg

instance (a ~ Fd) => QuickCmdArg [(a, IO Handle)] where
  quickCmdArg :: [(a, IO Handle)] -> Cmd -> Cmd
quickCmdArg = (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c. (a -> b -> c) -> b -> a -> c
flip) forall a. QuickCmdArg a => a -> Cmd -> Cmd
quickCmdArg

instance (a ~ Fd) => QuickCmdArg [(a, Handle)] where
  quickCmdArg :: [(a, Handle)] -> Cmd -> Cmd
quickCmdArg = (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c. (a -> b -> c) -> b -> a -> c
flip) forall a. QuickCmdArg a => a -> Cmd -> Cmd
quickCmdArg

instance QuickCmdArg ByteString where
  quickCmdArg :: ByteString -> Cmd -> Cmd
quickCmdArg = ByteString -> Cmd -> Cmd
passArg

instance (a ~ Fd, b ~ Cmd) => QuickCmdArg (a, (Handle -> IO b) -> IO b) where
  quickCmdArg :: (a, (Handle -> IO b) -> IO b) -> Cmd -> Cmd
quickCmdArg (a
new, (Handle -> IO b) -> IO b
getOld) Cmd
cmd = IO Cmd -> Cmd
unIOCmd forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Handle -> IO b) -> IO b
getOld forall a b. (a -> b) -> a -> b
$ \Handle
old -> do
    Fd
old <- Handle -> IO Fd
handleToFd Handle
old
    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ (Fd, Fd) -> Cmd -> Cmd
passFd (a
new, Fd
old) Cmd
cmd

instance (a ~ Fd) => QuickCmdArg (a, IO Handle) where
  quickCmdArg :: (a, IO Handle) -> Cmd -> Cmd
quickCmdArg (a
new, IO Handle
old) Cmd
cmd = IO Cmd -> Cmd
unIOCmd forall a b. (a -> b) -> a -> b
$ do
    Fd
old <- IO Handle
old forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Handle -> IO Fd
handleToFd
    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ (Fd, Fd) -> Cmd -> Cmd
passFd (a
new, Fd
old) Cmd
cmd

instance (a ~ Fd) => QuickCmdArg (a, Handle) where
  quickCmdArg :: (a, Handle) -> Cmd -> Cmd
quickCmdArg (a
new, Handle
old) Cmd
cmd = IO Cmd -> Cmd
unIOCmd forall a b. (a -> b) -> a -> b
$ do
    Fd
old <- Handle -> IO Fd
handleToFd Handle
old
    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ (Fd, Fd) -> Cmd -> Cmd
passFd (a
new, Fd
old) Cmd
cmd

instance QuickCmdArg (Cmd -> Cmd) where
  quickCmdArg :: (Cmd -> Cmd) -> Cmd -> Cmd
quickCmdArg = forall a. a -> a
id

instance {-# OVERLAPPABLE #-} (QuickCmdArg a, QuickCmd b) => QuickCmd (a -> b) where
  quickCmd :: Cmd -> a -> b
quickCmd Cmd
cmd a
arg = forall a. QuickCmd a => Cmd -> a
quickCmd forall a b. (a -> b) -> a -> b
$ forall a. QuickCmdArg a => a -> Cmd -> Cmd
quickCmdArg a
arg Cmd
cmd

instance (a ~ ()) => QuickCmd (IO a) where
  quickCmd :: Cmd -> IO a
quickCmd = Cmd -> IO ()
run

instance QuickCmd Cmd where
  quickCmd :: Cmd -> Cmd
quickCmd = forall a. a -> a
id

{- | >>> mq "cat" "/dev/null" (pipeArgIn 1 $ mq "cat" "/dev/null") <<< "somestr"

 The first argument is the path, and the subsequent arguments are 'QuickCmdArg'.
 At the end you will either have an @IO ()@ (synchronous execution) or 'Cmd' (which you can further use).
-}
mq ::
  (QuickCmd a, ToByteString b) =>
  -- | The path to the executable, uses PATH
  b ->
  -- | Either a 'Cmd', an @IO ()@, or a function that takes @Cmd -> Cmd@ , 'String' or 'ByteString'
  a
mq :: forall a b. (QuickCmd a, ToByteString b) => b -> a
mq b
path = forall a. QuickCmd a => Cmd -> a
quickCmd forall a b. (a -> b) -> a -> b
$ ByteString -> Cmd
makeCmd (forall a. ToByteString a => a -> ByteString
toByteString b
path)

{- | Pipe from the right command to the left command.
 Returns the left command modified.
-}
infixr 1 <|

(<|) :: QuickCmd a => Cmd -> Cmd -> a
<| :: forall a. QuickCmd a => Cmd -> Cmd -> a
(<|) Cmd
x Cmd
y = forall a. QuickCmd a => Cmd -> a
quickCmd forall a b. (a -> b) -> a -> b
$ Fd -> Fd -> Cmd -> Cmd -> Cmd
pipeIn Fd
1 Fd
0 Cmd
y Cmd
x

{- | Pipe from the right command's stderr to the left command.
 Returns the left command modified.
-}
infixr 1 <!|

(<!|) :: QuickCmd a => Cmd -> Cmd -> a
<!| :: forall a. QuickCmd a => Cmd -> Cmd -> a
(<!|) Cmd
x Cmd
y = forall a. QuickCmd a => Cmd -> a
quickCmd forall a b. (a -> b) -> a -> b
$ Fd -> Fd -> Cmd -> Cmd -> Cmd
pipeIn Fd
2 Fd
0 Cmd
y Cmd
x

{- | Pipe from the left command to the right command.
 Returns the left command modified.
-}
infixr 1 |>

(|>) :: QuickCmd a => Cmd -> Cmd -> a
|> :: forall a. QuickCmd a => Cmd -> Cmd -> a
(|>) Cmd
x Cmd
y = forall a. QuickCmd a => Cmd -> a
quickCmd forall a b. (a -> b) -> a -> b
$ Fd -> Fd -> Cmd -> Cmd -> Cmd
pipeOut Fd
0 Fd
1 Cmd
y Cmd
x

{- | Pipe from the left command's stderr to the right command.
 Returns the left command modified.
-}
infixr 1 |!>

(|!>) :: QuickCmd a => Cmd -> Cmd -> a
|!> :: forall a. QuickCmd a => Cmd -> Cmd -> a
(|!>) Cmd
x Cmd
y = forall a. QuickCmd a => Cmd -> a
quickCmd forall a b. (a -> b) -> a -> b
$ Fd -> Fd -> Cmd -> Cmd -> Cmd
pipeOut Fd
0 Fd
2 Cmd
y Cmd
x

-- | Pass a string as stdin.
infix 1 <<<

(<<<) :: (QuickCmd a, ToByteString b) => Cmd -> b -> a
<<< :: forall a b. (QuickCmd a, ToByteString b) => Cmd -> b -> a
(<<<) Cmd
cmd b
str = forall a. QuickCmd a => Cmd -> a
quickCmd forall a b. (a -> b) -> a -> b
$ Fd -> (Async ProcessStatus -> Handle -> IO ()) -> Cmd -> Cmd
pipeHIn Fd
0 (\Async ProcessStatus
_ Handle
h -> Handle -> ByteString -> IO ()
B.hPut Handle
h (forall a. ToByteString a => a -> ByteString
toByteString b
str) forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Handle -> IO ()
hClose Handle
h) Cmd
cmd

-- This function is pretty much never useful. If you want to handle the output
-- of the command, use `capture` or similar.
-- The problem is that it creates a new thread in the background, when
-- what we really want is to handle the output in the foreground, because
-- when our foreground is done executing, it will not wait for the background threads
-- to stop executing too.
---- | Handle the output from stdout.
-- infixl 1 >>>
--
-- (>>>) :: QuickCmd a => Cmd -> (ByteString -> IO ()) -> a
-- (>>>) cmd handler = quickCmd $ pipeHOut 1 (\_ h -> B.hGetContents h >>= handler) cmd

-- Disabled with same reason as for `>>>`.
---- | Handle the output from stderr.
----infixl 1 !>>>
----
----(!>>>) :: QuickCmd a => Cmd -> (ByteString -> IO ()) -> a
----(!>>>) cmd handler = quickCmd $ pipeHOut 2 (\_ h -> B.hGetContents h >>= handler) cmd

{- | Pass an argument of the form @\/proc\/self\/fd\/\<n\>@ to the process,
 where `n` is the reader end of a pipe which the passed string is written to.
-}
pipeArgStrIn :: ToByteString b => b -> Cmd -> Cmd
pipeArgStrIn :: forall b. ToByteString b => b -> Cmd -> Cmd
pipeArgStrIn b
str = (Async ProcessStatus -> Handle -> IO ()) -> Cmd -> Cmd
pipeArgHIn (\Async ProcessStatus
_ Handle
h -> Handle -> ByteString -> IO ()
B.hPut Handle
h (forall a. ToByteString a => a -> ByteString
toByteString b
str) forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Handle -> IO ()
hClose Handle
h)

-- Disabled with same reason as for `>>>`.
-- pipeArgStrOut :: (ByteString -> IO ()) -> Cmd -> Cmd
-- pipeArgStrOut handler = pipeArgHOut (\_ h -> B.hGetContents h >>= handler)

attachFinalizer :: IO () -> ByteString -> IO ByteString
attachFinalizer :: IO () -> ByteString -> IO ByteString
attachFinalizer IO ()
finalizer ByteString
str = [ByteString] -> ByteString
B.fromChunks forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [ByteString] -> IO [ByteString]
go (ByteString -> [ByteString]
B.toChunks ByteString
str)
  where
    go' :: [BS.ByteString] -> IO [BS.ByteString]
    go' :: [ByteString] -> IO [ByteString]
go' [] = IO ()
finalizer forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (f :: * -> *) a. Applicative f => a -> f a
pure []
    go' (ByteString
x : [ByteString]
xs) = (ByteString
x forall a. a -> [a] -> [a]
:) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [ByteString] -> IO [ByteString]
go [ByteString]
xs

    go :: [BS.ByteString] -> IO [BS.ByteString]
    go :: [ByteString] -> IO [ByteString]
go = forall a. IO a -> IO a
unsafeInterleaveIO forall b c a. (b -> c) -> (a -> b) -> a -> c
. [ByteString] -> IO [ByteString]
go'

{- | Capture the output of the fd of the command lazily.
 If the process exits with a non-zero exit code,
 reading from the bytestring will throw 'Procex.Process.CmdException'.
 Garbage collection will close the pipe.
-}
captureFdLazy :: Fd -> Cmd -> IO ByteString
captureFdLazy :: Fd -> Cmd -> IO ByteString
captureFdLazy Fd
fd Cmd
cmd = do
  (Async ProcessStatus
status, [Item [Handle]
h]) <- [Fd] -> Cmd -> IO (Async ProcessStatus, [Handle])
captureFdsAsHandles [Fd
fd] Cmd
cmd
  ByteString
out <- Handle -> IO ByteString
B.hGetContents Item [Handle]
h
  IO () -> ByteString -> IO ByteString
attachFinalizer (Async ProcessStatus -> IO ()
waitCmd Async ProcessStatus
status) ByteString
out

{- | Capture the output of the fd of the command lazily. Ignores process exit code.
 Garbage collection will close the pipe.
-}
captureFdLazyNoThrow :: Fd -> Cmd -> IO ByteString
captureFdLazyNoThrow :: Fd -> Cmd -> IO ByteString
captureFdLazyNoThrow Fd
fd Cmd
cmd = do
  (Async ProcessStatus
_, [Item [Handle]
h]) <- [Fd] -> Cmd -> IO (Async ProcessStatus, [Handle])
captureFdsAsHandles [Fd
fd] Cmd
cmd
  Handle -> IO ByteString
B.hGetContents Item [Handle]
h

-- | 'captureFdLazy' for stdout.
captureLazy :: Cmd -> IO ByteString
captureLazy :: Cmd -> IO ByteString
captureLazy = Fd -> Cmd -> IO ByteString
captureFdLazy Fd
1

-- | 'captureFdLazy' for stderr..
captureErrLazy :: Cmd -> IO ByteString
captureErrLazy :: Cmd -> IO ByteString
captureErrLazy = Fd -> Cmd -> IO ByteString
captureFdLazy Fd
2

-- | 'captureFdLazyNoThrow' for stdout.
captureLazyNoThrow :: Cmd -> IO ByteString
captureLazyNoThrow :: Cmd -> IO ByteString
captureLazyNoThrow = Fd -> Cmd -> IO ByteString
captureFdLazyNoThrow Fd
1

-- | 'captureFdLazyNoThrow' for stderr.
captureErrLazyNoThrow :: Cmd -> IO ByteString
captureErrLazyNoThrow :: Cmd -> IO ByteString
captureErrLazyNoThrow = Fd -> Cmd -> IO ByteString
captureFdLazyNoThrow Fd
2

captureFd' :: Fd -> Cmd -> IO (Async ProcessStatus, ByteString)
captureFd' :: Fd -> Cmd -> IO (Async ProcessStatus, ByteString)
captureFd' Fd
fd Cmd
cmd = do
  (Async ProcessStatus
status, [Item [Handle]
h]) <- [Fd] -> Cmd -> IO (Async ProcessStatus, [Handle])
captureFdsAsHandles [Fd
fd] Cmd
cmd
  !ByteString
out <- forall a. NFData a => a -> a
force forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Handle -> IO ByteString
B.hGetContents Item [Handle]
h
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (Async ProcessStatus
status, ByteString
out)

-- | Capture the output of the fd of the command strictly, err if the command exits with a non-zero exit code.
captureFd :: Fd -> Cmd -> IO ByteString
captureFd :: Fd -> Cmd -> IO ByteString
captureFd Fd
fd Cmd
cmd = do
  (Async ProcessStatus
status, ByteString
out) <- Fd -> Cmd -> IO (Async ProcessStatus, ByteString)
captureFd' Fd
fd Cmd
cmd
  Async ProcessStatus -> IO ()
waitCmd Async ProcessStatus
status
  forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
out

-- | Capture the output of the fd of the command strictly. Ignores process exit code.
captureFdNoThrow :: Fd -> Cmd -> IO ByteString
captureFdNoThrow :: Fd -> Cmd -> IO ByteString
captureFdNoThrow Fd
fd Cmd
cmd = do
  (Async ProcessStatus
_, ByteString
out) <- Fd -> Cmd -> IO (Async ProcessStatus, ByteString)
captureFd' Fd
fd Cmd
cmd
  forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
out

-- | 'captureFd' for stdout.
capture :: Cmd -> IO ByteString
capture :: Cmd -> IO ByteString
capture = Fd -> Cmd -> IO ByteString
captureFd Fd
1

-- | 'captureFdNoThrow' for stdout.
captureNoThrow :: Cmd -> IO ByteString
captureNoThrow :: Cmd -> IO ByteString
captureNoThrow = Fd -> Cmd -> IO ByteString
captureFdNoThrow Fd
1

-- | 'captureFd' for stderr.
captureErr :: Cmd -> IO ByteString
captureErr :: Cmd -> IO ByteString
captureErr = Fd -> Cmd -> IO ByteString
captureFd Fd
2

-- | 'captureFdNoThrow' for stderr.
captureErrNoThrow :: Cmd -> IO ByteString
captureErrNoThrow :: Cmd -> IO ByteString
captureErrNoThrow = Fd -> Cmd -> IO ByteString
captureFdNoThrow Fd
2