Comand line arguments as a StateT.

cargs.hs

{-# LANGUAGE TemplateHaskell, Rank2Types, NoMonomorphismRestriction #-}
import Data.List 
import Data.Maybe
import Data.Monoid
import Control.Monad.State 
import Control.Monad.Trans.State (state)
import Control.Monad.Identity
import Control.Monad
import Control.Monad.Writer 
import Control.Applicative

import qualified System.Environment as Env

import Control.Lens

{- ArgList element functions. -}
isShort :: String -> Bool
isShort a = a !! 0 == '-' && a !! 1 /= '-'

shortValue :: String -> String
shortValue = drop 1

matchShort :: Char -> String -> Bool
matchShort c arg = isShort arg && shortValue arg == [c]

charToShort :: Char -> String
charToShort c = '-' : [c]

isLong :: String -> Bool
isLong a = length a > 2 && dash 0 && dash 1 && not (dash 2)
    where dash i = a !! i == '-'

longValue :: String -> String
longValue = drop 2

matchLong :: String -> String -> Bool
matchLong str a = isLong a && longValue a == str

toLong :: String -> String
toLong str = "--" ++ str

data Argument = Long    String
              | Short   String -- Even a short arg, like -j12, may be more than
                               -- one Char
              | Regular String
              | Aliases [Argument] deriving (Eq)

isLongArg :: Argument -> Bool
isLongArg (Long _) = True
isLongArg _ = False

isShortArg :: Argument -> Bool
isShortArg (Short _) = True
isShortArg _ = False

isRegularArg :: Argument -> Bool
isRegularArg (Regular _) = True
isRegularArg _ = False

isAliases :: Argument -> Bool
isAliases (Aliases _) = True
isAliases _ = False

readArg :: Read r => Argument -> r
readArg (Long    l) = read l
readArg (Short   c) = read c
readArg (Regular r) = read r
readArg (Aliases args) = readArg (args !! 0) -- Shouldn't actually be called.

instance Show Argument where
    show (Long str)  = "--" ++ str
    show (Short c )  = '-' : c
    show (Regular r) = r
    show (Aliases as) = show $ map show as

short :: Char -> Argument
short c = Short [c]

{- Obtain any value data from an Argument.

    value (Long "--op=val") == Just "val"
    value (Short "-j12")    == Just "12" 
    value (Short "-x")      == Nothing
    value (Regular r)       == Just r
-}
value :: Argument -> Maybe String
value (Long l) = do i <- findIndex (=='=') l 
                    Just $ drop (i+1) l
value (Short c) = case drop 1 c of
    "" -> Nothing
    c' -> Just c'
value (Regular r) = Just r

readValue :: Read r => Argument -> Maybe r
readValue arg = value arg >>= withValue
    where
    withValue v = case reads v of [(r,"")] -> Just r
                                  _        -> Nothing

hasValue :: Argument -> Bool
hasValue arg = isJust (value arg)

cArg :: String -> Argument
cArg arg | isShort arg = Short (shortValue arg)
         | isLong  arg = Long  (longValue arg)
         | otherwise   = Regular arg

type ArgList = [Argument]

cArgs :: [String] -> ArgList
cArgs = map cArg

aliases :: [String] -> Argument
aliases = Aliases . cArgs

match :: Argument -> Argument -> Bool
match (Aliases as) (Aliases bs) = False -- Aliases don't match with each other.
match (Aliases as) b = matchAny as b
match a (Aliases bs) = matchAny bs a
match a b = isPrefixOf a' b' || isPrefixOf b' a'
    where a' = show a
          b' = show b

matchAny :: [Argument] -> Argument -> Bool
matchAny [] _ = False
matchAny (a:as) b = match a b || matchAny as b

matchLongShort :: String -> Char -> Argument -> Bool
matchLongShort l s a = matchAny [Long l, Short [s]] a

{- ArgT : An argument list state transformer. -

An Arg modifies a state (type: ArgList) and returns a (value,arglist) pair. -}
type ArgT m r = StateT ArgList m r
type Arg r = ArgT Identity r

{- State and StateT specializations. -}
runArg  = runState
evalArg = evalState
execArg = execState

runArg_  a = runState  a []
evalArg_ a = evalState a []
execArg_ a = execState a []

mapArg  = mapState
withArg = withState

runArgT  = runStateT
evalArgT = evalStateT
execArgT = execStateT
mapArgT  = mapStateT
withArgT = withStateT

runArgT_  a = runStateT  a []
evalArgT_ a = evalStateT a []
execArgT_ a = execStateT a []

{-- Monad State specializations. --}
arg  = state
argT = StateT

putBack :: Monad m => Argument -> ArgT m ()
putBack arg = modify ((:) arg)

-- const modify
modify_ ::Monad m => ArgList -> ArgT m ()
modify_ = modify . const

{- IO Utilities -}
getArgs :: IO (ArgList)
getArgs = fmap cArgs Env.getArgs 

type ArgIO r = ArgT IO r

printArg :: Show r => ArgT IO r -> ArgList -> IO ()
printArg a al = runArgT a al >>= print

{- Maybe utilities -}
{- Convert an ArgT Maybe to an ArgT m

Leaving the Maybe monad to a different one can be tricky because its state will
be lost if one just calls runArgT. The original ArgT must be run, and the state
of its return must be inserted into the new ArgT.
-}
maybeArg :: Monad m
         => r -> (a -> r) -> ArgT Maybe a -> ArgT m r
maybeArg r f ma = do al <- get
                     maybe (return r) update (runArgT ma al)
    where update (a,al) = put al >> return (f a)

fromMaybeArg :: Monad m
             => ArgT m r -> (a -> ArgT m r) -> ArgT Maybe a 
             -> ArgT m r
fromMaybeArg r f ma = do al <- get 
                         maybe al (runArgT ma al)
    where maybe al Nothing        = r
          maybe al (Just (a,al')) = put al' >> f a

{- Examples. -}
{- Gobble and return first argument. -

runArg gobbleFirst [a,b] == (Just a, [b]) 
runArg gobbleFirst []    == (Nothing, []) -}
gobbleFirst :: ArgT Maybe Argument
gobbleFirst = get >>= gobble
    where gobble []     = lift Nothing
          gobble (a:al) = put al >> return a

gobbleRegular :: ArgT Maybe Argument
gobbleRegular = get >>= gobble []
    where gobble _   []     = lift Nothing
          gobble acc (Regular r : al) = do put (acc ++ al)
                                           return (Regular r)
          gobble acc (a:al) = gobble (acc++[a]) al
                                                

{- Show each argument. -

This function has no real purpose except to demonstrate how to access each
argument. 

runArg showArgs lst == (unwords lst, lst) -}
showArgs :: Monad m => ArgT m String
showArgs = get >>= return . unwords . map show 

deleteAt :: Int -> [a] -> [a]
deleteAt 0 (x:xs) = xs
deleteAt _ []     = []
deleteAt n (x:xs) = x : deleteAt (n-1) xs

findArgIndex :: (Argument -> Bool) -> ArgT Maybe Int
findArgIndex pred = get >>= maybe (lift Nothing) return . findIndex pred

gobbleAt :: Monad m => Int -> ArgT m Argument
gobbleAt i = do a <- get >>= return . (!! i)
                modify (deleteAt i)
                return a

gobbleBy :: (Argument -> Bool) -> ArgT Maybe Argument
gobbleBy pred = findArgIndex pred >>= gobbleAt 

{- Gobble an Argument.

runArg (gobbleArg (Long l)) [toLong l] == (True,[])
runArg (gobbleArg (Long k)) [toLong k] == (False,[k]) -}
gobbleArg :: Monad m => Argument -> ArgT m Bool
gobbleArg a = maybeArg False (const True) (gobbleBy $ match a)

type ArgP r = ArgT Maybe (Maybe r)

badParse :: Read r => ArgP r
badParse = fail "parse error"

notFound :: Read r => ArgP r
notFound = lift (Just Nothing)

returnValue :: Read r => r -> ArgP r
returnValue = lift . Just . Just

parsedOk :: Maybe (Maybe r, ArgList) -> Bool
parsedOk (Just _) = True
parsedOk _ = False

found :: Maybe (Maybe r, ArgList) -> Bool
found (Just (Just _, _)) = True
found _ = False

parsed :: Maybe (Maybe r, ArgList) -> r
parsed (Just (Just x,_)) = x

unprocessed :: Maybe (Maybe r, ArgList) -> ArgList
unprocessed (Just (_,al)) = al

-- TODO: Should this take a third argument, (a -> r), instead of using id?
fromArgP :: Monad m => r -> m r -> ArgP r -> ArgT m r
fromArgP r a = fromMaybeArg (lift a) 
                            (return . maybe r id)

-- Workaround for argP <|> argP producing the wrong value.
(<||>) :: ArgP r -> ArgP r -> ArgP r
a <||> b = do al <- get
              case runArgT a al of
                Just (Nothing,_)   -> b
                Just (Just x, al') -> put al' >> return (Just x)
                Nothing            -> b

{- Gobble an argument/value pair.

Search for a given handle and return its value. 
If the Argument does not have a value, return the value of the next Argument if
it has the type Regular.

A result of Nothing means a parse error. The state is abandoned.
    ex: let j = gobbleValue (short 'j') :: Arg Int 
        runArgT j [cArg "-jn"] == Nothing -- Can't parse Int from 'n'.

A result of Just (Nothing,args) means args does not contain handle.
    ex: runArgT j [cArg "-k"] == Just (Nothing,[-k])

A result of Just (Just x) represents a succesful parse.
    ex: evalArgT j [cArg "-j4"] == Just 4

TODO: This behaviour makes (gobbleValue x <|> gobbleValue y) work improperly in
      that if gobbleValue x returns Just (Nothing,al), gobbleValue y never
      runs. 
-}

argPValue :: Read r => Argument -> ArgP r
argPValue arg = 
    maybe (lift Nothing) (return . Just) (readValue arg)

gobbleValue :: Read r => Argument -> ArgP r
gobbleValue handle =
    do al <- get
       let m = runArgT (findArgIndex $ match handle) al
       case m of Nothing    -> return Nothing
                 Just (i,_) -> gobble i
    where 
    gobble i = do a <- gobbleAt i 
                  argPValue a <|> tryNext
        where tryNext = do a <- gobbleAt i
                           guard (isRegularArg a)
                           argPValue a

gobbleShort :: Monad m => Char -> ArgT m Bool
gobbleShort = gobbleArg . short

-- Gobble: Long version.
gobbleLong :: Monad m => String -> ArgT m Bool
gobbleLong = gobbleArg . Long

{-- Testing. --}
testList :: ArgList
testList = cArgs ["-j3","-x", "-y", "-z", "--long", "--help"]

showEnvArgs :: IO (String)
showEnvArgs = getArgs >>= evalArgT showArgs

sumArgs' :: ArgT (Writer (Sum Int)) ()
sumArgs' = do fromMaybeArg done rec gobbleFirst
    where 
    done = return ()

    -- Try to read the values of Regulars only.
    rec (Regular n) = tellArg (Regular n) (readValue $ Regular n)

    -- Put other args back in the state after we've finished.
    rec a = rec' a 
    rec' a = do sumArgs'
                putBack a

    -- Either keep n and tell the writer, or put it back.
    tellArg a (Just n) = do sumArgs'
                            lift $ tell (Sum n)
    tellArg a Nothing = rec' a
                       

sumArgs :: Monad m => ArgT m Int
sumArgs = do al <- get
             let w = runArgT sumArgs' al
             let (((),al'),sum) = runWriter w
             put al' >> return (getSum sum)

-- Like the make option, -jn
jobs :: ArgT Maybe Int
jobs = fromArgP 1       -- Default if -j not specified.
                Nothing -- Failled to parse.
                        -- Found -j, return n.
                (gobbleValue $ aliases ["-j","--jobs"])

{- TestOptions -- A less trivial example using a TestOptions record. -}
data TestOptions = TestOptions 
    { _tstOpA :: Bool -- -a
    , _tstOpB :: Bool -- -b
    , _tstOpC :: Bool -- -c
    } deriving (Show)

testOptionsDefaults = return $ TestOptions False False False

makeLenses ''TestOptions

-- Make an ArgT from a lens.
testOptionArg :: Monad m
              => Char 
              -> Lens TestOptions TestOptions Bool Bool
              -> TestOptions -> ArgT m TestOptions
testOptionArg c opLens ops = liftM set (gobbleShort c)
    where set value = (opLens .~ value) ops

-- Link a, b, and c to their respective lenses.
testOptionsA = testOptionArg 'a' tstOpA
testOptionsB = testOptionArg 'b' tstOpB
testOptionsC = testOptionArg 'c' tstOpC

-- Test options by setting the default options as the accumulator
-- and testing for A, then B, then C.
testOptions = testOptionsDefaults >>= testOptionsA
                                  >>= testOptionsB
                                  >>= testOptionsC 

-- Produce a truth table of results.
runTestOptions :: IO ()
runTestOptions = let
    argsArgs = [ cArgs ["-a", "-b", "-c"]                
               , cArgs ["-x", "-b", "-c"]
               , cArgs ["-a", "-x", "-c"]
               , cArgs ["-a", "-b", "-x"]
               , cArgs ["-x", "-x", "-c"]
               , cArgs ["-x", "-b", "-x"]
               , cArgs ["-a", "-x", "-x"]
               , cArgs ["-x", "-x", "-x"]
               ]
    in mapM_ (printArg testOptions) argsArgs