Ticket #3430: CmdArgs.hs

File CmdArgs.hs, 8.2 KB (added by NeilMitchell, 6 years ago)
Line 
1{-# LANGUAGE ScopedTypeVariables, DeriveDataTypeable #-}
2{-|
3    Simple command line argument handling
4-}
5
6module System.Console.CmdArgs where
7
8import Prelude hiding (catch)
9import System.IO.Unsafe
10import Data.Dynamic
11import Data.Data
12import Data.List
13import Data.Maybe
14import Data.IORef
15import System.Environment
16import Control.Monad
17import Control.Exception
18import System.Exit
19import System.FilePath
20import Data.Char
21import Control.Monad.State
22
23
24---------------------------------------------------------------------
25-- DEFAULTS
26
27class Default a where
28    def :: a
29
30instance Default Bool where def = False
31instance Default [a] where def = []
32
33---------------------------------------------------------------------
34-- STATE MANAGEMENT
35
36{-# NOINLINE verbosity #-}
37verbosity :: IORef Int -- 0 = quiet, 1 = normal, 2 = verbose
38verbosity = unsafePerformIO $ newIORef 1
39
40isQuiet, isNormal, isLoud :: IO Bool
41isQuiet = return True
42isNormal = fmap (>=1) $ readIORef verbosity
43isLoud = fmap (>=2) $ readIORef verbosity
44
45
46---------------------------------------------------------------------
47-- STATE MANAGEMENT
48
49{-# NOINLINE info #-}
50info :: IORef [Info]
51info = unsafePerformIO $ newIORef []
52
53(&) :: forall a . a -> Info -> a
54(&) x i = unsafePerformIO $ do
55    modifyIORef info (i:)
56    return x
57
58collect :: IO [Info]
59collect = do
60    x <- readIORef info
61    writeIORef info []
62    return x
63
64---------------------------------------------------------------------
65-- USER INTERFACE
66
67type Flag = [Info]
68type Mode = [Info]
69
70data Info
71    = FldName String -- the record name
72    | FldType TypeRep -- the field type
73    | FldEmpty Dynamic
74    | FldTyp String
75    | FldText String
76    | FldFlag String
77    | FldArgs
78    | HelpSuffix [String]
79      deriving Show
80
81isFldArgs FldArgs{} = True; isFldArgs _ = False
82
83isFlag = not . any isFldArgs
84
85
86-- | A default argument if none is specified
87empty :: Typeable a => a -> Info
88empty = FldEmpty . toDyn
89
90-- | The type of the argument
91typ :: String -> Info
92typ = FldTyp
93
94-- | Descriptive text for the option
95text :: String -> Info
96text = FldText
97
98-- | Flags which work
99flag :: String -> Info
100flag = FldFlag
101
102-- | Where to put the non-flag arguments
103args :: Info
104args = FldArgs
105
106
107typDir, typFile :: Info
108typFile = typ "FILE"
109typDir = typ "DIR"
110
111
112helpSuffix :: [String] -> Info
113helpSuffix = HelpSuffix
114
115
116---------------------------------------------------------------------
117-- MAIN DRIVERS
118
119cmdArgs :: Data a => String -> a -> IO a
120cmdArgs short x = do
121    evaluate x
122    mode <- collect
123    ref <- newIORef (constrFields $ toConstr x, [])
124    x <- flip gmapM x $ \i -> do
125        res <- evaluate i
126        info <- collect
127        let typ = typeOf i
128        modifyIORef ref $ \(fld:flds, xs) ->
129            if knownType typ
130            then (flds, (FldName fld:FldType typ:info):xs)
131            else error $ "Can't handle a type of " ++ fld
132        return res
133    flags <- fmap (reverse . snd) $ readIORef ref
134    flags <- return $ assignLong flags
135    flags <- return $ assignShort flags
136
137    args <- concatMap expandShort `fmap` getArgs
138    when (any (`elem` args) ["-?","--help"]) $ do
139        showHelp short mode flags
140        exitSuccess
141    when (any (`elem` args) ["-V","--version"]) $ do
142        putStrLn short
143        exitSuccess
144    process flags args x
145
146
147cmdArgsMode :: Data a => String -> [a] -> IO a
148cmdArgsMode = error "todo"
149
150
151---------------------------------------------------------------------
152-- UTILITIES
153
154expandShort :: String -> [String]
155expandShort ('-':x:xs) | x /= '-' && xs /= "" = ['-',x] : expandShort ('-':xs)
156expandShort x = [x]
157
158
159assignLong :: [Flag] -> [Flag]
160assignLong = map f
161    where f xs = [FldFlag $ flagName x | FldName x <- xs, isFlag xs] ++ xs
162
163
164reservedShort = "?Vvq"
165
166assignShort :: [Flag] -> [Flag]
167assignShort xs = zipWith (++) flags xs
168    where
169        seen = [y | x <- xs, FldFlag [y] <- x]
170        guesses = map guess xs
171        dupes = nub $ concat guesses \\ nub (concat guesses)
172        flags = [[FldFlag [i] | i <- g, i `notElem` (seen++dupes)] | g <- guesses]
173
174        -- guess at a possible short flag
175        guess ys = if [() | FldFlag [_] <- ys] /= [] then [] else take 1 [x | FldFlag (x:_) <- ys]
176
177
178flagName :: String -> String
179flagName xs | "_" `isSuffixOf` xs = flagName $ init xs
180flagName xs = [if x == '_' then '-' else x | x <- xs]
181
182
183errorIO :: String -> IO a
184errorIO x = putStrLn x >> exitFailure
185
186
187setField :: Data a => a -> String -> (Dynamic -> Dynamic) -> a
188setField x name v = flip evalState (constrFields $ toConstr x) $ flip gmapM x $ \i -> do
189    n:ns <- get
190    put ns
191    return $ if n == name then fromDyn (v $ toDyn i) i else i
192
193
194---------------------------------------------------------------------
195-- HELP INFORMATION
196
197showHelp :: String -> Mode -> [Flag] -> IO ()
198showHelp short mode flags = do
199    x <- getProgName
200    let ty = head $ [y | x <- flags, any isFldArgs x, FldTyp y <- x] ++ ["FILE"]
201    showBlock $
202        Left short :
203        Left "" :
204        Left ("  " ++ map toLower (takeBaseName x) ++ " [FLAG] ["++ty++"]") :
205        Left "" :
206        Right ("-?","--help","Show usage information") :
207        Right ("-V","--version","Show version information") :
208        Right ("-v","--verbose","Higher verbosity") :
209        Right ("-q","--quiet","Lower verbosity") :
210        concatMap (map Right . showArg) flags ++
211        concat [map Left $ "":xs | HelpSuffix xs <- mode]
212
213
214showBlock :: [Either String (String,String,String)] -> IO ()
215showBlock xs = putStr $ unlines $ map f xs
216    where f (Left x) = x
217          f (Right (a,b,c)) = "  " ++ pad an a ++ pad bn b ++ c
218
219          (as,bs,_) = unzip3 [x | Right x <- xs]
220          an = maximum $ map length as
221          bn = maximum $ map length bs
222          pad n x = x ++ replicate (n - length x + 2) ' '
223
224
225showArg :: Flag -> [(String,String,String)]
226showArg xs =
227    [(unwords (map ("-"++) short) ++ if null short then "" else val True
228     ,unwords (map ("--"++) long) ++ val False
229     ,unwords [x | FldText x <- xs])
230    | isFlag xs]
231    where
232        (short,long) = partition ((==) 1 . length) [x | FldFlag x <- xs]
233        val b = if or [boolType x | FldType x <- xs] then ""
234                else [' ' | not opt && b] ++ ['['|opt] ++ ['='|not b] ++ head ([x | FldTyp x <- xs] ++ ["VALUE"]) ++ [']'|opt]
235        opt = [] /= [() | FldEmpty x <- xs]
236
237
238---------------------------------------------------------------------
239-- PROCESS FLAGS
240
241process :: Data a => [Flag] -> [String] -> a -> IO a
242process flags [] a = return a
243process flags (x:xs) a
244    | x `elem` ["-v","--verbose"] = writeIORef verbosity 2 >> process flags xs a
245    | x `elem` ["-q","--quiet"  ] = writeIORef verbosity 0 >> process flags xs a
246
247process flags xs a = f flags
248    where
249        f [] = errorIO $ "Unhandled flag: " ++ head xs
250        f (y:ys) = case processOne y xs a of
251            Nothing -> f ys
252            Just v -> do (a,xs) <- v; process flags xs a
253
254
255processOne :: Data a => Flag -> [String] -> a -> Maybe (IO (a, [String]))
256processOne flag o@(('-':x):xs) val | or [n == a | FldFlag n <- flag] = Just $
257    case typ of
258        _ | boolType typ -> do
259            when (b /= "") $ errorIO $ "Error on flag " ++ head o ++ ", " ++ a ++ " does not take an argument"
260            return (setValue $ const $ toDyn True, xs)
261        _ -> error "Unhandled type"
262    where
263        (a,b) = break (== '=') (if "-" `isPrefixOf` x then tail x else x)
264        typ = head [x | FldType x <- flag] 
265        setValue = setField val (head [x | FldName x <- flag])
266
267processOne flag o@(x:xs) val | any isFldArgs flag = error "argument here"
268
269processOne _ _ _ = Nothing
270
271
272knownType :: TypeRep -> Bool
273knownType x = basicType x || maybe False basicType (listType x)
274    where basicType x = boolType x || isJust (readType x)
275
276
277listType :: TypeRep -> Maybe TypeRep
278listType x | a == typeRepTyCon (typeOf "") = Just $ head b
279    where (a,b) = splitTyConApp x
280listType _ = Nothing
281
282
283boolType :: TypeRep -> Bool
284boolType x = x == typeOf True 
285
286
287readType :: TypeRep -> Maybe (String -> Dynamic)
288readType x
289    | x == typeOf "" = Just toDyn
290    | x == typeOf (0 :: Int) = Just $ toDyn . (read :: String -> Int)
291    | x == typeOf (0 :: Integer) = Just $ toDyn . (read :: String -> Integer)
292    | x == typeOf (0 :: Float) = Just $ toDyn . (read :: String -> Float)
293    | x == typeOf (0 :: Double) = Just $ toDyn . (read :: String -> Double)
294    | otherwise = Nothing