{-# LANGUAGE OverloadedStrings, QuasiQuotes, BangPatterns, ScopedTypeVariables #-}
module Main where
import Data.Text (Text)
import Control.Applicative
import Control.Monad.Reader
import qualified Data.Map.Strict as M
import qualified Data.ByteString.Lazy as L
import Data.ByteString.Lazy as BL hiding (map, intersperse, zip, concat, putStrLn)
import qualified Data.ByteString.Lazy.Char8 as L8
import System.Environment (getArgs)
import Data.Aeson
import Data.Monoid
import Data.List (intersperse, zip)
import qualified Data.Attoparsec.Text as AT
import Data.Attoparsec.Lazy as Atto hiding (Result)
import Data.Attoparsec.ByteString.Char8 (endOfLine, sepBy)
import qualified Data.HashMap.Lazy as HM
import qualified Data.Vector as V
import Data.Scientific (Scientific, floatingOrInteger)
import qualified Data.Text as T
import qualified Data.Text.Lazy.IO as TL
import qualified Options.Applicative as O
import Control.Monad (when)
import System.Exit
import Data.String.QQ
import System.Process (readProcess, readProcessWithExitCode)
data Options = Options {
preFilter :: Maybe String
, filterProgram :: String
, fieldsToFilter :: String
, debugKeyPaths :: Bool
} deriving Show
parseOpts :: O.Parser Options
parseOpts = Options
<$> ((Just <$> O.strOption (O.short 'p' <> O.long "prefilter" <> O.metavar "PREFILTER" <> O.help "This filter prevent main filter from being applied if exit code 0")) <|> pure Nothing)
<*> O.argument O.str (O.metavar "FILTERPROG" <> O.help "External filter")
<*> O.argument O.str (O.metavar "FIELDS" <> O.help "JSON keypath expressions")
<*> O.switch (O.long "debug" <> O.help "Debug keypaths")
opts = O.info (O.helper <*> parseOpts)
(O.fullDesc
<> O.progDesc [s|Filter JSON object fields through shell.
On STDIN provide an input stream of newline-separated JSON objects. |]
<> O.header "jsonextfilter"
<> O.footer "See https://github.com/danchoi/jsonextfilter for more information.")
main = do
Options preFilter filterProg expr debugKeyPaths <- O.execParser opts
let (prog:args) = words filterProg
x <- BL.getContents
let xs :: [Value]
xs = decodeStream x
ks = parseKeyPath $ T.pack expr
ks' :: [[Key]]
ks' = [k | KeyPath k <- ks]
when debugKeyPaths $ do
putStrLn $ "FilterProg: " ++ prog ++ " " ++ show args
putStrLn $ "PreFilterProg: " ++ show (fmap words preFilter)
Prelude.putStrLn $ "Key Paths: " ++ show ks
exitSuccess
-- transform JSON
let bashFilter = io prog args preFilter
xs' <- mapM (runFilterOnPaths bashFilter ks') xs
mapM_ (L8.putStrLn . encode) xs'
io :: String -> [String] -> Maybe String -> Value -> IO Value
io prog args preFilter v =
case v of
String v' -> do
go <- runPreFilter v'
if go
then do
res <- readProcess prog args (T.unpack v')
return . String . T.pack $ res
else return v -- no op
_ -> return v -- no op
where
runPreFilter v =
case preFilter of
Just preFilter' -> do
let preProg:pArgs = words preFilter'
(exitcode, _, _ ) <- readProcessWithExitCode preProg pArgs (T.unpack v)
case exitcode of
ExitSuccess -> return True
ExitFailure _ -> return False
_ -> return True
------------------------------------------------------------------------
runFilterOnPaths :: (Value -> IO Value) -> [[Key]] -> Value -> IO Value
runFilterOnPaths ioFilter ks v =
foldM
(\acc kp ->
case acc of
(Object acc') -> do
v' <- runReaderT (runFilterOnPath [] acc) (FilterEnv kp ioFilter)
case v' of
(Object hm) -> return . Object $ HM.union hm acc'
x -> error $ "Expected Object, but got " ++ show x
x -> error $ "Expected Object, but got " ++ show x
) v ks
data FilterEnv = FilterEnv { targetKeyPath :: [Key]
, filterProg :: (Value -> IO Value)
}
runFilterOnPath :: [Key] -> Value -> ReaderT FilterEnv IO Value
runFilterOnPath k v = do
-- liftIO $ putStrLn $ "runFilterPath " ++ show k
targetKeyPath' <- asks targetKeyPath
bashFilter' <- asks filterProg
if (k == targetKeyPath')
then liftIO $ bashFilter' v
else go k v
where
go :: [Key] -> Value -> ReaderT FilterEnv IO Value
go _ x@(String _) = return x
go _ Null = return Null
go _ x@(Number _) = return x
go _ x@(Bool _) = return x
go _ x@(Array _) = return x -- no effect on Arrays
go ks x@(Object hm) = do
let pairs = HM.toList hm
pairs' <- mapM (\(k,v) -> (,) <$> pure k <*> runFilterOnPath (ks <> [Key k]) v) pairs
return . Object . HM.fromList $ pairs'
------------------------------------------------------------------------
-- decode JSON object stream
decodeStream :: (FromJSON a) => BL.ByteString -> [a]
decodeStream bs = case decodeWith json bs of
(Just x, xs) | xs == mempty -> [x]
(Just x, xs) -> x:(decodeStream xs)
(Nothing, _) -> []
decodeWith :: (FromJSON a) => Parser Value -> BL.ByteString -> (Maybe a, BL.ByteString)
decodeWith p s =
case Atto.parse p s of
Atto.Done r v -> f v r
Atto.Fail _ _ _ -> (Nothing, mempty)
where f v' r = (\x -> case x of
Success a -> (Just a, r)
_ -> (Nothing, r)) $ fromJSON v'
------------------------------------------------------------------------
-- JSON parsing and data extraction
data KeyPath = KeyPath [Key] deriving Show
data Key = Key Text | Index Int deriving (Eq, Show)
parseKeyPath :: Text -> [KeyPath]
parseKeyPath s = case AT.parseOnly pKeyPaths s of
Left err -> error $ "Parse error " ++ err
Right res -> res
spaces = many1 AT.space
pKeyPaths :: AT.Parser [KeyPath]
pKeyPaths = pKeyPath `AT.sepBy` spaces
pKeyPath :: AT.Parser KeyPath
pKeyPath = KeyPath
<$> (AT.sepBy1 pKeyOrIndex (AT.takeWhile1 $ AT.inClass ".["))
pKeyOrIndex :: AT.Parser Key
pKeyOrIndex = pIndex <|> pKey
pKey = Key <$> AT.takeWhile1 (AT.notInClass " .[:")
pIndex = Index <$> AT.decimal <* AT.char ']'