summaryrefslogtreecommitdiff
path: root/Coalpit/IO.hs
blob: b0a074a41a4d8f205f8fd219e2bb2029cb274858 (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
{- |
Module      :  Coalpit.IO
Description :  Helper IO functions
Maintainer  :  defanor <defanor@uberspace.net>
Stability   :  unstable
Portability :  non-portable (uses GHC extensions)

These are basic utility functions for pipes-based IO.

An example:

@
\{\-\# LANGUAGE DeriveGeneric, DeriveAnyClass \#\-\}
import GHC.Generics
import "Coalpit"

data Output = Foo Double | Bar deriving (Generic, 'Coalpit')
data Args = Args { arg1 :: Maybe Int, arg2 :: Double }
  deriving (Generic, 'Coalpit')

main :: IO ()
main = 'pipe' $ \a i -> pure $ Foo $ maybe (arg2 a) fromIntegral (arg1 a) + i
@

-}

{-# LANGUAGE RankNTypes, ScopedTypeVariables #-}

module Coalpit.IO ( run, run', handleErrors
                  , pipe, consumer, producer) where

import Data.Proxy (Proxy(..))
import System.Environment (getProgName, getArgs)
import Control.Monad.IO.Class (MonadIO, liftIO)
import System.Exit (die)
import System.IO (hPutStrLn, stderr)
import Pipes ((>->), Pipe, yield, await, lift, runEffect)
import Control.Monad (mapM_, forever)
import qualified Pipes.Prelude as PP
import Coalpit.Core (Coalpit, fromArgs, defOpt, usage)
import Coalpit.DSV (readDSV, showDSV)

-- | Runs a given action on each 'Left' value, embedding that action's
-- result into the data stream.
handleErrors :: MonadIO m => (e -> m [a]) -> Pipe (Either e a) a m ()
handleErrors e = forever $ do
  v <- await
  case v of
    Left err -> do
      vs <- lift $ e err
      mapM_ yield vs
    Right x -> yield x

-- | Runs a given 'Pipe' between input producer and output consumer.
-- Prints an error and usage instructions if it fails to parse the
-- arguments, and passes the input through 'handleErrors'.
run :: forall m a i o. (MonadIO m, Coalpit a, Coalpit i, Coalpit o)
    => (String -> m [i])
    -- ^ An action to run on error (see 'handleErrors').
    -> (a -> Pipe i o m ())
    -- ^ Main function.
    -> m ()
run e f = do
  pn <- liftIO getProgName
  let u = Prelude.concat ["Usage: ", pn, " ", usage defOpt (Proxy :: Proxy a)]
  args <- liftIO getArgs
  a <- either (liftIO . die . (++ u)) pure $ fromArgs defOpt args
  runEffect $
    PP.stdinLn
    >-> PP.map (readDSV defOpt)
    >-> handleErrors e
    >-> f a
    >-> PP.map (showDSV defOpt)
    >-> PP.stdoutLn

-- | Same as 'run', but just prints errors into 'stderr'.
run' :: forall m a i o. (MonadIO m, Coalpit a, Coalpit i, Coalpit o)
         => (a -> Pipe i o m ())
         -- ^ Main function.
         -> m ()
run' = run (\e -> liftIO $ hPutStrLn stderr e >> pure [])

-- | Similar to 'interact' in the approach, a wrapper around 'run''.
pipe :: (MonadIO m, Coalpit a, Coalpit i, Coalpit o)
     => (a -> i -> m o) -> m ()
pipe f = run' $ \a -> forever $ await >>= lift . f a >>= yield

-- | A version of 'pipe' that doesn't 'await' any input.
producer :: forall m a o. (MonadIO m, Coalpit a, Coalpit o)
         => (a -> m o) -> m ()
producer f = run' $ \a -> forever $ lift (f a) >>= yield :: Pipe () o m ()

-- | A version of 'pipe' that doesn't 'yield' any output.
consumer :: forall m a i. (MonadIO m, Coalpit a, Coalpit i)
         => (a -> i -> m ()) -> m ()
consumer f = run' $ \a -> forever $ await >>= lift . f a :: Pipe i () m ()