Enable color only for terminal output, add manual options

Changelog: Use colors by default only on terminal, add
  `--color`/`--no-color` options to select manually.

2 files changed, 27 insertions, 9 deletions

diff --git a/src/Main.hs b/src/Main.hs
index f296a12..594de0c 100644
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -14,6 +14,8 @@ import System.Exit
 import System.FilePath
 import System.FilePath.Glob
 import System.IO
+import System.Posix.Terminal
+import System.Posix.Types
 
 import Config
 import Output
@@ -28,6 +30,7 @@ data CmdlineOptions = CmdlineOptions
     { optTest :: TestOptions
     , optRepeat :: Int
     , optVerbose :: Bool
+    , optColor :: Maybe Bool
     , optShowHelp :: Bool
     , optShowVersion :: Bool
     }
@@ -37,6 +40,7 @@ defaultCmdlineOptions = CmdlineOptions
     { optTest = defaultTestOptions
     , optRepeat = 1
     , optVerbose = False
+    , optColor = Nothing
     , optShowHelp = False
     , optShowVersion = False
     }
@@ -52,6 +56,12 @@ options =
     , Option ['v'] ["verbose"]
         (NoArg (\opts -> opts { optVerbose = True }))
         "show output of processes and successful tests"
+    , Option [] [ "color" ]
+        (NoArg (\opts -> opts { optColor = Just True }))
+        "always use colors for output (default when stdout is tty)"
+    , Option [] [ "no-color" ]
+        (NoArg (\opts -> opts { optColor = Just False }))
+        "never use colors for output (default when stdout is not a tty)"
     , Option ['t'] ["timeout"]
         (ReqArg (\str -> to $ \opts -> case readMaybe str of
                                             Just timeout -> opts { optTimeout = timeout }
@@ -133,7 +143,10 @@ main = do
 
     when (null files) $ fail $ "No test files"
 
-    out <- startOutput $ optVerbose opts
+    useColor <- case optColor opts of
+        Just use -> return use
+        Nothing -> queryTerminal (Fd 1)
+    out <- startOutput (optVerbose opts) useColor
 
     tests <- forM files $ \(path, mbTestName) -> do
         fileTests <- parseTestFile path
diff --git a/src/Output.hs b/src/Output.hs
index 96bb965..135e6e0 100644
--- a/src/Output.hs
+++ b/src/Output.hs
@@ -27,6 +27,7 @@ data Output = Output
 
 data OutputConfig = OutputConfig
     { outVerbose :: Bool
+    , outUseColor :: Bool
     }
 
 data OutputState = OutputState
@@ -50,10 +51,10 @@ class MonadIO m => MonadOutput m where
 instance MonadIO m => MonadOutput (ReaderT Output m) where
     getOutput = ask
 
-startOutput :: Bool -> IO Output
-startOutput verbose = Output
+startOutput :: Bool -> Bool -> IO Output
+startOutput outVerbose outUseColor = Output
     <$> newMVar OutputState { outPrint = TL.putStrLn, outHistory = emptyHistory }
-    <*> pure OutputConfig { outVerbose = verbose }
+    <*> pure OutputConfig { .. }
 
 outColor :: OutputType -> Text
 outColor OutputChildStdout = T.pack "0"
@@ -97,11 +98,15 @@ outLine :: MonadOutput m => OutputType -> Maybe Text -> Text -> m ()
 outLine otype prompt line = ioWithOutput $ \out ->
     when (outVerbose (outConfig out) || printWhenQuiet otype) $ do
         withMVar (outState out) $ \st -> do
-            outPrint st $ TL.fromChunks
-                [ T.pack "\ESC[", outColor otype, T.pack "m"
-                , maybe "" (<> outSign otype <> outArr otype <> " ") prompt
-                , line
-                , T.pack "\ESC[0m"
+            outPrint st $ TL.fromChunks $ concat
+                [ if outUseColor (outConfig out)
+                    then [ T.pack "\ESC[", outColor otype, T.pack "m" ]
+                    else []
+                , [ maybe "" (<> outSign otype <> outArr otype <> " ") prompt ]
+                , [ line ]
+                , if outUseColor (outConfig out)
+                    then [ T.pack "\ESC[0m" ]
+                    else []
                 ]
 
 outPromptGetLine :: MonadOutput m => Text -> m (Maybe Text)