add global option `--color`

1 files changed, 38 insertions, 12 deletions

diff --git a/app/Main.hs b/app/Main.hs
index 1617c86..e7bf0bc 100644
--- a/app/Main.hs
+++ b/app/Main.hs
@@ -377,16 +377,24 @@ import TreeGrepper.Match qualified as G
 import Prelude hiding (id)
 
 data Options = Options
-  { internalTags :: Bool,
-    command :: Command
+  { command :: Command,
+    internalTags :: Bool,
+    colorize :: Color
   }
   deriving (Show)
 
+data Color
+  = Always
+  | Auto
+  | Never
+  deriving (Show, Eq)
+
 options :: O.Parser Options
 options =
   Options
-    <$> internalTagsFlag
-    <*> cmd
+    <$> cmd
+    <*> internalTagsFlag
+    <*> colorOption
 
 internalTagsFlag :: O.Parser Bool
 internalTagsFlag =
@@ -395,6 +403,21 @@ internalTagsFlag =
         <> O.help "Whether to display internal tags."
     )
 
+colorOption :: O.Parser Color
+colorOption =
+  O.option
+    (O.maybeReader parse)
+    ( O.long "color"
+        <> O.short 'c'
+        <> O.help "Wether to colorize output. (Default: auto)"
+        <> O.value Auto
+    )
+  where
+    parse "auto" = pure Auto
+    parse "always" = pure Always
+    parse "never" = pure Never
+    parse _ = Nothing
+
 data Command
   = List
       { files :: [String],
@@ -462,9 +485,9 @@ die s = do
 main :: IO ()
 main = do
   O.execParser (O.info (options <**> O.helper) O.idm) >>= \case
-    Options {command = List {sort, filters, files}} -> do
+    Options {colorize, command = List {sort, filters, files}} -> do
       issues <- listIssues sort filters files
-      putDoc . P.vsep $
+      putDoc colorize . P.vsep $
         map
           ( \issue ->
               let title = P.annotate P.bold $ P.pretty issue.title
@@ -495,12 +518,12 @@ main = do
                     ]
           )
           issues
-    Options {command = Show {id, width}} -> do
+    Options {colorize, command = Show {id, width}} -> do
       issues <- listIssues [] [] []
       case find ((==) (Just id) . I.id) issues of
         Nothing -> die (printf "no issue with id `%s'\n" id)
         Just issue -> do
-          putDoc $
+          putDoc colorize $
             P.annotate (P.color P.Green) $
               P.pretty $
                 issue.file
@@ -536,7 +559,7 @@ main = do
                       String.fromString (printf "mdcat --columns %d --local" width')
                 )
             )
-          putDoc $
+          putDoc colorize $
             P.pretty $
               "\n@file "
                 ++ issue.file
@@ -544,8 +567,8 @@ main = do
                 ++ show issue.start.row
                 ++ "\n"
 
-putDoc :: P.Doc P.AnsiStyle -> IO ()
-putDoc doc = do
+putDoc :: Color -> P.Doc P.AnsiStyle -> IO ()
+putDoc colorize doc = do
   isTty <- (== 1) <$> c_isatty 1
   columns <- fmap Terminal.width <$> Terminal.size
   P.renderIO stdout
@@ -553,7 +576,10 @@ putDoc doc = do
       P.defaultLayoutOptions
         { P.layoutPageWidth = maybe P.Unbounded (\n -> P.AvailablePerLine n 1) columns
         }
-    $ (if isTty then (\x -> x) else P.unAnnotate)
+    $ ( if colorize == Always || (colorize == Auto && isTty)
+          then (\x -> x)
+          else P.unAnnotate
+      )
     $ doc
 
 foreign import ccall "unistd.h isatty" c_isatty :: Int -> IO Int