{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE OverloadedStrings    #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-| Useful helpers to style and color text with ANSI escape sequences.
-}
module System.Console.Pretty
( Color(..) , Pretty(..) , Section(..) , Style(..)
, supportsPretty)
where

import qualified Data.Char          as C
import           Data.Monoid        ((<>))
import qualified Data.Text          as T
import           GHC.IO.Handle      (Handle)
import           System.Environment (lookupEnv)
import           System.IO          (hIsTerminalDevice, stdout)

---------------------------------------------------------------------------------
-- TYPES

-- | A section to be colored, either foreground or background.
data Section = Foreground | Background

-- | Colors for an ANSI terminal
data Color = Black | Red | Green | Yellow | Blue | Magenta | Cyan | White | Default
  deriving (Int -> Color
Color -> Int
Color -> [Color]
Color -> Color
Color -> Color -> [Color]
Color -> Color -> Color -> [Color]
(Color -> Color)
-> (Color -> Color)
-> (Int -> Color)
-> (Color -> Int)
-> (Color -> [Color])
-> (Color -> Color -> [Color])
-> (Color -> Color -> [Color])
-> (Color -> Color -> Color -> [Color])
-> Enum Color
forall a.
(a -> a)
-> (a -> a)
-> (Int -> a)
-> (a -> Int)
-> (a -> [a])
-> (a -> a -> [a])
-> (a -> a -> [a])
-> (a -> a -> a -> [a])
-> Enum a
enumFromThenTo :: Color -> Color -> Color -> [Color]
$cenumFromThenTo :: Color -> Color -> Color -> [Color]
enumFromTo :: Color -> Color -> [Color]
$cenumFromTo :: Color -> Color -> [Color]
enumFromThen :: Color -> Color -> [Color]
$cenumFromThen :: Color -> Color -> [Color]
enumFrom :: Color -> [Color]
$cenumFrom :: Color -> [Color]
fromEnum :: Color -> Int
$cfromEnum :: Color -> Int
toEnum :: Int -> Color
$ctoEnum :: Int -> Color
pred :: Color -> Color
$cpred :: Color -> Color
succ :: Color -> Color
$csucc :: Color -> Color
Enum)

-- | SGR paramaters, aka text styles for an ANSI terminal
data Style
  = Normal | Bold | Faint | Italic
  | Underline | SlowBlink | ColoredNormal | Reverse
  deriving (Int -> Style
Style -> Int
Style -> [Style]
Style -> Style
Style -> Style -> [Style]
Style -> Style -> Style -> [Style]
(Style -> Style)
-> (Style -> Style)
-> (Int -> Style)
-> (Style -> Int)
-> (Style -> [Style])
-> (Style -> Style -> [Style])
-> (Style -> Style -> [Style])
-> (Style -> Style -> Style -> [Style])
-> Enum Style
forall a.
(a -> a)
-> (a -> a)
-> (Int -> a)
-> (a -> Int)
-> (a -> [a])
-> (a -> a -> [a])
-> (a -> a -> [a])
-> (a -> a -> a -> [a])
-> Enum a
enumFromThenTo :: Style -> Style -> Style -> [Style]
$cenumFromThenTo :: Style -> Style -> Style -> [Style]
enumFromTo :: Style -> Style -> [Style]
$cenumFromTo :: Style -> Style -> [Style]
enumFromThen :: Style -> Style -> [Style]
$cenumFromThen :: Style -> Style -> [Style]
enumFrom :: Style -> [Style]
$cenumFrom :: Style -> [Style]
fromEnum :: Style -> Int
$cfromEnum :: Style -> Int
toEnum :: Int -> Style
$ctoEnum :: Int -> Style
pred :: Style -> Style
$cpred :: Style -> Style
succ :: Style -> Style
$csucc :: Style -> Style
Enum)


---------------------------------------------------------------------------------
-- CLASS

-- | A class to color and style
class Pretty a where
  -- | Helper to set foreground color
  color :: Color -> a -> a
  color = Section -> Color -> a -> a
forall a. Pretty a => Section -> Color -> a -> a
colorize Section
Foreground
  -- | Helper to set background color
  bgColor :: Color -> a -> a
  bgColor = Section -> Color -> a -> a
forall a. Pretty a => Section -> Color -> a -> a
colorize Section
Background
  -- | Set the color of the (fg/bg) with the color
  colorize :: Section -> Color -> a -> a
  -- | Set the style
  style :: Style -> a -> a

---------------------------------------------------------------------------------
-- TEXT

-- | Instance of `Pretty` for `T.Text`
instance Pretty T.Text where
  colorize :: Section -> Color -> Text -> Text
colorize Section
section Color
col Text
str =
    Text
"\x1b[" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<>                                  -- escape code
    Text
sectionNum Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<>                               -- bg/foreground
    (Char -> Text
T.singleton (Char -> Text) -> Char -> Text
forall a b. (a -> b) -> a -> b
$ Int -> Char
C.intToDigit (Int -> Char) -> Int -> Char
forall a b. (a -> b) -> a -> b
$ Color -> Int
forall a. Enum a => a -> Int
fromEnum Color
col) -- color code
    Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"m" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<>                                   -- delim
    Text
str Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<>                                      -- inner string
    Text
"\x1b[0m"                                   -- reset
    where
      sectionNum :: T.Text
      sectionNum :: Text
sectionNum = case Section
section of
        Section
Foreground -> Text
"9"
        Section
Background -> Text
"4"

  style :: Style -> Text -> Text
style Style
sty Text
str =
    Text
"\x1b[" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<>                                  -- escape code
    (Char -> Text
T.singleton (Char -> Text) -> Char -> Text
forall a b. (a -> b) -> a -> b
$ Int -> Char
C.intToDigit (Int -> Char) -> Int -> Char
forall a b. (a -> b) -> a -> b
$ Style -> Int
forall a. Enum a => a -> Int
fromEnum Style
sty) -- style
    Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"m" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<>                                   -- delim
    Text
str Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<>                                      -- inner string
    Text
"\x1b[0m"                                   -- reset

---------------------------------------------------------------------------------
-- STRING

-- | Instance of `Pretty` for `String`
instance Pretty String where
  colorize :: Section -> Color -> String -> String
colorize Section
section Color
col String
str =
    String
"\x1b[" String -> String -> String
forall a. Semigroup a => a -> a -> a
<>          -- escape code
    String
sectionNum String -> String -> String
forall a. Semigroup a => a -> a -> a
<>       -- bg/foreground
    Int -> String
forall a. Show a => a -> String
show (Color -> Int
forall a. Enum a => a -> Int
fromEnum Color
col) -- color code
    String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"m" String -> String -> String
forall a. Semigroup a => a -> a -> a
<>           -- delim
    String
str String -> String -> String
forall a. Semigroup a => a -> a -> a
<>              -- inner string
    String
"\x1b[0m"           -- reset
    where
      sectionNum :: String
      sectionNum :: String
sectionNum = case Section
section of
        Section
Foreground -> String
"9"
        Section
Background -> String
"4"

  style :: Style -> String -> String
style Style
sty String
str =
    String
"\x1b[" String -> String -> String
forall a. Semigroup a => a -> a -> a
<>             -- escape code
    Int -> String
forall a. Show a => a -> String
show (Style -> Int
forall a. Enum a => a -> Int
fromEnum Style
sty)    -- style
    String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"m" String -> String -> String
forall a. Semigroup a => a -> a -> a
<>              -- delim
    String
str String -> String -> String
forall a. Semigroup a => a -> a -> a
<>                 -- string
    String
"\x1b[0m"              -- reset


---------------------------------------------------------------------------------
-- SUPPORTED CHECK

-- | Whether or not the current terminal supports pretty-terminal
supportsPretty :: IO Bool
supportsPretty :: IO Bool
supportsPretty =
  Handle -> IO Bool
hSupportsANSI Handle
stdout
  where
    -- | Use heuristics to determine whether the functions defined in this
    -- package will work with a given handle.
    --
    -- The current implementation checks that the handle is a terminal, and
    -- that the @TERM@ environment variable doesn't say @dumb@ (whcih is what
    -- Emacs sets for its own terminal).
    hSupportsANSI :: Handle -> IO Bool
    -- Borrowed from an HSpec patch by Simon Hengel
    -- (https://github.com/hspec/hspec/commit/d932f03317e0e2bd08c85b23903fb8616ae642bd)
    hSupportsANSI :: Handle -> IO Bool
hSupportsANSI Handle
h = Bool -> Bool -> Bool
(&&) (Bool -> Bool -> Bool) -> IO Bool -> IO (Bool -> Bool)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Handle -> IO Bool
hIsTerminalDevice Handle
h IO (Bool -> Bool) -> IO Bool -> IO Bool
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Bool -> Bool
not (Bool -> Bool) -> IO Bool -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO Bool
isDumb)
      where
        isDumb :: IO Bool
isDumb = (Maybe String -> Maybe String -> Bool
forall a. Eq a => a -> a -> Bool
== String -> Maybe String
forall a. a -> Maybe a
Just String
"dumb") (Maybe String -> Bool) -> IO (Maybe String) -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> IO (Maybe String)
lookupEnv String
"TERM"