lib: useColor, colorOption helpers usable anywhere

hledger-lib/Hledger/Reports/ReportOptions.hs

@@ -48,15 +48,12 @@ where
 
 import Control.Applicative ((<|>))
 import Data.List.Extra (nubSort)
-import Data.Maybe (fromMaybe, isJust, mapMaybe)
+import Data.Maybe (fromMaybe, mapMaybe)
 import qualified Data.Text as T
 import Data.Time.Calendar (Day, addDays)
 import Data.Default (Default(..))
 import Safe (headMay, lastDef, lastMay, maximumMay)
 
-import System.Console.ANSI (hSupportsANSIColor)
-import System.Environment (lookupEnv)
-import System.IO (stdout)
 import Text.Megaparsec.Custom
 
 import Hledger.Data
@@ -181,11 +178,8 @@ defreportopts = ReportOpts
 rawOptsToReportOpts :: RawOpts -> IO ReportOpts
 rawOptsToReportOpts rawopts = do
     d <- getCurrentDay
-    no_color <- isJust <$> lookupEnv "NO_COLOR"
-    supports_color <- hSupportsANSIColor stdout
 
-    let colorflag    = stringopt "color" rawopts
-        formatstring = T.pack <$> maybestringopt "format" rawopts
+    let formatstring = T.pack <$> maybestringopt "format" rawopts
         querystring  = map T.pack $ listofstringopt "args" rawopts  -- doesn't handle an arg like "" right
         (costing, valuation) = valuationTypeFromRawOpts rawopts
 
@@ -222,10 +216,7 @@ rawOptsToReportOpts rawopts = do
           ,percent_     = boolopt "percent" rawopts
           ,invert_      = boolopt "invert" rawopts
           ,pretty_tables_ = boolopt "pretty-tables" rawopts
-          ,color_       = and [not no_color
-                              ,not $ colorflag `elem` ["never","no"]
-                              ,colorflag `elem` ["always","yes"] || supports_color
-                              ]
+          ,color_       = useColor -- a lower-level helper
           ,forecast_    = forecastPeriodFromRawOpts d rawopts
           ,transpose_   = boolopt "transpose" rawopts
           }

hledger-lib/Hledger/Utils/Debug.hs

@@ -94,7 +94,7 @@ module Hledger.Utils.Debug (
   ,traceParse
   ,dbgparse
   ,module Debug.Trace
-)
+,useColor)
 where
 
 import           Control.Monad (when)
@@ -105,12 +105,15 @@ import qualified Data.Text.Lazy as TL
 import           Debug.Trace
 import           Hledger.Utils.Parse
 import           Safe (readDef)
-import           System.Environment (getArgs)
+import           System.Environment (getArgs, lookupEnv)
 import           System.Exit
 import           System.IO.Unsafe (unsafePerformIO)
 import           Text.Megaparsec
 import           Text.Printf
 import           Text.Pretty.Simple  -- (defaultOutputOptionsDarkBg, OutputOptions(..), pShowOpt, pPrintOpt)
+import Data.Maybe (isJust)
+import System.Console.ANSI (hSupportsANSIColor)
+import System.IO (stdout)
 
 prettyopts = 
     defaultOutputOptionsDarkBg
@@ -139,6 +142,49 @@ ptrace = traceWith pshow
 traceWith :: Show a => (a -> String) -> a -> a
 traceWith f a = trace (f a) a
 
+-- | Check the IO environment to see if ANSI colour codes should be used in output.
+-- This is done using unsafePerformIO so it can be used anywhere, eg in 
+-- low-level debug utilities, which should be ok since we are just reading.
+-- (When running code in GHCI, this module must be reloaded to see a change.)
+-- The logic is: use color if 
+-- a NO_COLOR environment variable is not defined
+-- and the program was not started with --color=no|never
+-- and stdout supports ANSI color, or the program was started with --color=yes|always.
+-- {-# OPTIONS_GHC -fno-cse #-}
+-- {-# NOINLINE useColor #-}
+useColor :: Bool
+useColor = unsafePerformIO $ do
+  no_color       <- isJust <$> lookupEnv "NO_COLOR"
+  supports_color <- hSupportsANSIColor stdout
+  let coloroption = colorOption
+  return $ and [
+     not no_color
+    ,not $ coloroption `elem` ["never","no"]
+    ,coloroption `elem` ["always","yes"] || supports_color
+    ]
+
+-- | Read the value of the --color or --colour command line option provided at program startup
+-- using unsafePerformIO. If this option was not provided, returns the empty string.
+-- (When running code in GHCI, this module must be reloaded to see a change.)
+-- {-# OPTIONS_GHC -fno-cse #-}
+-- {-# NOINLINE colorOption #-}
+colorOption :: String
+colorOption = 
+  -- similar to debugLevel
+  let args = unsafePerformIO getArgs in
+  case snd $ break (=="--color") args of
+    "--color":v:_ -> v
+    _ ->
+      case take 1 $ filter ("--color=" `isPrefixOf`) args of
+        ['-':'-':'c':'o':'l':'o':'r':'=':v] -> v
+        _ ->
+          case snd $ break (=="--colour") args of
+            "--colour":v:_ -> v
+            _ ->
+              case take 1 $ filter ("--colour=" `isPrefixOf`) args of
+                ['-':'-':'c':'o':'l':'o':'u':'r':'=':v] -> v
+                _ -> ""
+
 -- | Global debug level, which controls the verbosity of debug output
 -- on the console. The default is 0 meaning no debug output. The
 -- @--debug@ command line flag sets it to 1, or @--debug=N@ sets it to