nptools-0.6.0: git-prompt.hs

{-# LANGUAGE RecordWildCards #-}
import System.Console.ANSI
import System.Environment
import System.Exit
import System.IO.Error (catchIOError)
import HSH
import Control.Monad
import Control.Applicative
import Data.Char
import Data.Maybe
import Debug.Trace (trace)

-- TODO: show stash information

-- tell Zsh to not count the contents for the length
noLength :: String -> String
noLength = ("%{"++) . (++"%}")

color :: Color -> String
color c
  = setSGRCode [ SetConsoleIntensity BoldIntensity
               , SetColor Foreground Dull c]

-- I didn't found how to generate sgr0 from Sys.Console.ANSI
sgr0 :: String
sgr0 = "\SI"

data Colors = Colors
  { blue, magenta, green, red, noColor :: String }

noColors :: Colors
noColors =
  Colors{blue     = ""
        ,magenta  = ""
        ,green    = ""
        ,red      = ""
        ,noColor  = ""
        }

mapColors :: (String -> String) -> Colors -> Colors
mapColors f Colors{..} =
  Colors{blue     = f blue
        ,magenta  = f magenta
        ,green    = f green
        ,red      = f red
        ,noColor  = f noColor
        }

colors :: Colors
colors = Colors{blue     = color Blue
               ,magenta  = color Magenta
               ,green    = color Green
               ,red      = color Red
               ,noColor  = setSGRCode [] ++ sgr0}

statsInfo :: Colors -> IO String
statsInfo Colors{..} = do
  gitstat <-  fmap lines . run $ "git status 2> /dev/null"
         -|-  egrep "(# On branch|# Untracked|# Changes|# Your branch|# and have)"
  let f msg   = listToMaybe . map digits . egrep msg $ gitstat
      g msg   = fmap words . listToMaybe . egrep msg $ gitstat
      ahead   = f "# Your branch is ahead of "
      behind  = f "# Your branch is behind "
      diverged = parseDiverged =<< g "# and have"
      branch = fromMaybe (red ++ "ERR") (parseOnBranch =<< g "# On branch")
  return . concat $
    ["±"
    ,blue,":"
    ,branch
    ,green
    ,['!' | "# Changes to be committed:"        `elem` gitstat]
    ,['?' | "# Untracked files:"                `elem` gitstat
         || "# Changes not staged for commit:"  `elem` gitstat]
    ,maybe "" ((blue++":"++green++"-")++) behind
    ,maybe "" ((blue++":"++green++"+")++) ahead
    ,maybe "" (\(x,y) -> blue++":"++red++"+"++x++"-"++y) diverged
    ,noColor]
  where
    digits  = reverse . takeWhile isDigit . dropWhile (not . isDigit) . reverse
    parseDiverged ["#", "and", "have", x, "and", y, "different", "commit(s)", "each,", "respectively."] = Just (x, y)
    parseDiverged _ = trace "git-prompt: git status message seems to have changed (diverged branches)" Nothing
    parseOnBranch ["#", "On", "branch", x] = Just (magenta ++ x)
    parseOnBranch _ = Nothing

isBareRepo :: IO Bool
isBareRepo = readbool =<< run "git config core.bare"
  where readbool "false\n"  = return False
        readbool "true\n"   = return True
        readbool s          = fail ("Unexpected result: " ++ show s)

cond :: a -> a -> Bool -> a
cond x y b = if b then x else y

main :: IO ()
main = do
  args <- getArgs
  ref <- head . lines <$> run "git symbolic-ref HEAD 2> /dev/null"
  when (null ref) exitFailure
  let myColors
        | "--no-color" `elem` args = noColors
        | "--zsh-no-length" `elem` args = mapColors noLength colors
        | otherwise = colors
      Colors{..} = myColors
  stats <- catchIOError (statsInfo myColors) . const $
             ((blue++":")++) . cond (green ++ "bare") (red ++ "ERR")
               <$> isBareRepo
  putStrLn stats