-
Notifications
You must be signed in to change notification settings - Fork 844
/
Args.hs
98 lines (89 loc) · 3.53 KB
/
Args.hs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
{-# LANGUAGE OverloadedStrings #-}
-- | Parsing argument-like things.
module Data.Attoparsec.Args
( EscapingMode(..)
, argsParser
, parseArgs
, withInterpreterArgs
) where
import Control.Applicative
import Data.Attoparsec.Text ((<?>))
import qualified Data.Attoparsec.Text as P
import Data.Attoparsec.Types (Parser)
import Data.ByteString (ByteString)
import qualified Data.ByteString as S
import Data.Conduit
import qualified Data.Conduit.Binary as CB
import qualified Data.Conduit.List as CL
import Data.Text (Text)
import Data.Text.Encoding (decodeUtf8')
import System.Directory (doesFileExist)
import System.Environment (getArgs, withArgs)
import System.IO (IOMode (ReadMode), withBinaryFile)
-- | Mode for parsing escape characters.
data EscapingMode
= Escaping
| NoEscaping
deriving (Show,Eq,Enum)
-- | Parse arguments using 'argsParser'.
parseArgs :: EscapingMode -> Text -> Either String [String]
parseArgs mode t = P.parseOnly (argsParser mode) t
-- | A basic argument parser. It supports space-separated text, and
-- string quotation with identity escaping: \x -> x.
argsParser :: EscapingMode -> Parser Text [String]
argsParser mode = many (P.skipSpace *> (quoted <|> unquoted)) <*
P.skipSpace <* (P.endOfInput <?> "unterminated string")
where
unquoted = P.many1 naked
quoted = P.char '"' *> string <* P.char '"'
string = many (case mode of
Escaping -> escaped <|> nonquote
NoEscaping -> nonquote)
escaped = P.char '\\' *> P.anyChar
nonquote = P.satisfy (not . (=='"'))
naked = P.satisfy (not . flip elem ("\" " :: String))
-- | Use 'withArgs' on result of 'getInterpreterArgs'.
withInterpreterArgs :: String -> ([String] -> Bool -> IO a) -> IO a
withInterpreterArgs progName inner = do
(args, isInterpreter) <- getInterpreterArgs progName
withArgs args $ inner args isInterpreter
-- | Check if command-line looks like it's being used as a script interpreter,
-- and if so look for a @-- progName ...@ comment that contains additional
-- arguments.
getInterpreterArgs :: String -> IO ([String], Bool)
getInterpreterArgs progName = do
args0 <- getArgs
case args0 of
(x:_) -> do
isFile <- doesFileExist x
if isFile
then do
margs <-
withBinaryFile x ReadMode $ \h ->
CB.sourceHandle h
$= CB.lines
$= CL.map killCR
$$ sinkInterpreterArgs progName
return $ case margs of
Nothing -> (args0, True)
Just args -> (args ++ "--" : args0, True)
else return (args0, False)
_ -> return (args0, False)
where
killCR bs
| S.null bs || S.last bs /= 13 = bs
| otherwise = S.init bs
sinkInterpreterArgs :: Monad m => String -> Sink ByteString m (Maybe [String])
sinkInterpreterArgs progName =
await >>= maybe (return Nothing) checkShebang
where
checkShebang bs
| "#!" `S.isPrefixOf` bs = fmap (maybe Nothing parseArgs') await
| otherwise = return (parseArgs' bs)
parseArgs' bs =
case decodeUtf8' bs of
Left _ -> Nothing
Right t ->
case P.parseOnly (argsParser Escaping) t of
Right ("--":progName':rest) | progName' == progName -> Just rest
_ -> Nothing