group values of list of key values pairs

grouper.hs

{-# LANGUAGE RecordWildCards #-}
import Control.Arrow (second)
import Control.Applicative ((<$>), (<*>))
import Data.List (groupBy, genericLength)
import System.Console.GetOpt (ArgDescr(..), ArgOrder(..))
import System.Console.GetOpt (OptDescr(..), getOpt, usageInfo)
import System.IO (hPutStrLn, stderr)
import System.Environment (getArgs)
import System.Exit (exitFailure)

type Separator = Char
type Method = String

data Options = Options {
      sep    :: Separator
    , method :: Method
}

data Flag = HelpOpt | SeparatorOpt Separator | MethodOpt Method
    deriving (Show)

options :: [OptDescr Flag]
options =
    [ Option ['h'] ["help"]
        (NoArg  HelpOpt)                           "show help"
    , Option ['s'] ["separator"]
        (ReqArg (SeparatorOpt . head) "separator") "key-value separator"
    , Option ['m'] ["method"]
        (ReqArg MethodOpt "avg/sum")               "aggregating method"
    ]

methods :: (Fractional a, Real a) => [(Method, [a] -> a)]
methods = [("sum", sum), ("avg", avg)]

main = do
  mbOptions <- compileOpts
  case mbOptions of
    Just ctx -> interact (go ctx)
    Nothing  -> showHelp

go :: Options -> String -> String
go Options{..} = unlines

          -- combine key with val
          . map (\(a,b) -> a ++ [sep] ++ b)

          -- aggregate values
          . map group

          -- group keys
          . groupBy (\(a,_) (b,_) -> a == b)

          -- tuple per line
          . map (second (drop 1) . break (== sep)) -- remove head

          . lines

  where group :: [(String, String)] -> (String, String)
        group xs@((k,_):_) = (k, show $ combine $ map readSnd xs)

        combine = maybe (fail $ "Unknown method " ++ method) id
                        (lookup method methods)

-- from http://stackoverflow.com/a/2377067/1080523
avg :: (Real a, Fractional b) => [a] -> b
avg xs = realToFrac (sum xs) / genericLength xs

-- reads the second value. key is passed for pretty useful error message
readSnd :: Read a => (String, String) -> a
readSnd kv@(k, v) = case reads v of
              [(x, "")] -> x
              _ -> error $ "Cannot parse " ++ show kv

-- options parsing
compileOpts :: IO (Maybe Options)
compileOpts = do
    args <- getArgs
    case getOpt RequireOrder options args of
        (o, _, []) -> return $ case find o HelpOpt of
                                 Just _ -> Nothing
                                 Nothing -> getOptions o
        (_, _, e)  -> hPutStrLn stderr ("unknown options " ++ concat e) >>
                      exitFailure

  where getOptions opts = Options <$> (head <$> find opts (SeparatorOpt ';'))
                                  <*> find opts (MethodOpt "sum")

        -- all very hacky
        find ((SeparatorOpt x):xs) (SeparatorOpt _) = Just [x]
        find []                    (SeparatorOpt x) = Just [x]
        find ((MethodOpt x):xs) (MethodOpt _) = Just x
        find []                 (MethodOpt x) = Just x
        find (_:xs) opt = find xs opt 
        find _ HelpOpt = Just "h"

showHelp :: IO ()
showHelp = putStrLn (usageInfo "grouper by rethab" options)