JQ.hs

{-# LANGUAGE FlexibleContexts #-}
module Data.JQ where

import Control.Applicative ( Applicative, (<$>), (<*>) )
import Control.Monad ( (>=>) )
import Control.Monad.Reader.Class
import qualified Data.Aeson as JSON
import qualified Data.Aeson.Types as JSON
import qualified Data.HashMap.Strict as HM
import Data.List ( foldl1' )
import Data.Text ( Text )
import qualified Data.Text as T
import qualified Data.Vector as V

-- $setup
-- >>> :set -XOverloadedStrings
-- >>> import Data.Maybe

data JQArrayFilterExpr = JQArrayFilterKey Text
                       | JQArrayFilterAt Int
                       | JQArrayFilterFromTo (Maybe Int) (Maybe Int)
                         deriving Eq

instance Show JQArrayFilterExpr where
    show expr = "[" ++ show' expr ++ "]" where
        show' (JQArrayFilterKey key) = "\"" ++ T.unpack key ++ "\""
        show' (JQArrayFilterAt ix) = show ix
        show' (JQArrayFilterFromTo from to) = maybe "" show from ++ ":" ++ maybe "" show to

data JQObjectElem = JQObjectElem (Either Text JQFilter) (Either Text JQFilter)
                  | JQObjectShortcutElem Text
                    deriving Eq

instance Show JQObjectElem where
    show (JQObjectElem key value) = either T.unpack show key ++ ":" ++ either T.unpack show value
    show (JQObjectShortcutElem object) = T.unpack object

data JQFilter = JQDotFilter
              | JQFieldFilter JQFilter Text Bool
              | JQArrayFilter JQFilter JQArrayFilterExpr Bool
              | JQAllFilter JQFilter Bool
              | JQCommaFilter JQFilter JQFilter
              | JQBarFilter JQFilter JQFilter
              | JQArray [JQFilter]
              | JQObject [JQObjectElem]
                deriving Eq

instance Show JQFilter where
    show JQDotFilter = "."
    show (JQFieldFilter jq key p) = show jq ++ T.unpack key ++ [ '?' | p ]
    show (JQArrayFilter jq expr p) = shows jq "[" ++ show expr ++ "]" ++ [ '?' | p ]
    show (JQAllFilter q p) = shows q "[]" ++ [ '?' | p ]
    show (JQCommaFilter a b) = show a ++ "," ++ show b
    show (JQBarFilter a b) = show a ++ show b
    show (JQArray jqs) = "[" ++ foldl1' (\a b -> a ++ "," ++ b) (map show jqs) ++ "]"
    show (JQObject elems) = "{" ++ foldl1' (\a b -> a ++ "," ++ b) (map show elems) ++ "}"

data JQError = JQE
instance Show JQError where
    show _ = "JQE"

-- | jq filter
--
-- http://stedolan.github.io/jq/manual/#example1
-- >>> query JQDotFilter (JSON.toJSON "Hello, world!")
-- Right [String "Hello, world!"]
--
-- http://stedolan.github.io/jq/manual/#example2
-- >>> query (JQFieldFilter JQDotFilter "foo" False) (fromJust $ JSON.decode "{\"foo\": 42, \"bar\": \"less interesting data\"}")
-- Right [Number 42.0]
-- >>> query (JQFieldFilter JQDotFilter "foo" False) (fromJust $ JSON.decode "{\"notfoo\": true, \"alsonotfoo\": false}")
-- Right [Null]
-- >>> query (JQArrayFilter JQDotFilter (JQArrayFilterKey "foo") False) (fromJust $ JSON.decode "{\"foo\": 42}")
-- Right [Number 42.0]
--
-- http://stedolan.github.io/jq/manual/#example3
-- >>> query (JQFieldFilter JQDotFilter "foo" True ) (fromJust $ JSON.decode "{\"foo\": 42, \"bar\": \"less interesting data\"}")
-- Right [Number 42.0]
-- >>> query (JQFieldFilter JQDotFilter "foo" True) (fromJust $ JSON.decode "{\"notfoo\": true, \"alsonotfoo\": false}")
-- Right [Null]
-- >>> query (JQArrayFilter JQDotFilter (JQArrayFilterKey "foo") True) (fromJust $ JSON.decode "{\"foo\": 42}")
-- Right [Number 42.0]
-- >>> query (JQArrayFilter JQDotFilter (JQArrayFilterKey "foo") True) (fromJust $ JSON.decode "[1,2]")
-- Right []
--
-- http://stedolan.github.io/jq/manual/#example4
-- >>> query (JQArrayFilter JQDotFilter (JQArrayFilterAt 0) False) (fromJust $ JSON.decode "[{\"name\":\"JSON\", \"good\":true}, {\"name\":\"XML\", \"good\":false}]")
-- Right [Object (fromList [("name",String "JSON"),("good",Bool True)])]
-- >>> query (JQArrayFilter JQDotFilter (JQArrayFilterAt 2) False) (fromJust $ JSON.decode "[{\"name\":\"JSON\", \"good\":true}, {\"name\":\"XML\", \"good\":false}]")
-- Right [Null]
-- >>> query (JQArrayFilter JQDotFilter (JQArrayFilterFromTo (Just 2) (Just 4)) False) (fromJust $ JSON.decode "[\"a\",\"b\",\"c\",\"d\",\"e\"]")
-- Right [String "c",String "d"]
-- >>> query (JQArrayFilter JQDotFilter (JQArrayFilterFromTo (Just 2) (Just 4)) False) (JSON.toJSON "abcdefghi")
-- Right [String "cd"]
-- >>> query (JQArrayFilter JQDotFilter (JQArrayFilterFromTo Nothing (Just 3)) False) (fromJust $ JSON.decode "[\"a\",\"b\",\"c\",\"d\",\"e\"]")
-- Right [String "a",String "b",String "c"]
-- >>> query (JQArrayFilter JQDotFilter (JQArrayFilterFromTo (Just (-2)) Nothing) False) (fromJust $ JSON.decode "[\"a\",\"b\",\"c\",\"d\",\"e\"]")
-- Right [String "d",String "e"]
--
-- http://stedolan.github.io/jq/manual/#example5
-- >>> query (JQAllFilter JQDotFilter False) (fromJust $ JSON.decode "[{\"name\":\"JSON\", \"good\":true}, {\"name\":\"XML\", \"good\":false}]")
-- Right [Object (fromList [("name",String "JSON"),("good",Bool True)]),Object (fromList [("name",String "XML"),("good",Bool False)])]
-- >>> query (JQAllFilter JQDotFilter False) (fromJust $ JSON.decode "[]")
-- Right []
-- >>> query (JQAllFilter JQDotFilter False) (fromJust $ JSON.decode "{\"a\": 1, \"b\": 1}")
-- Right [Number 1.0,Number 1.0]
--
-- http://stedolan.github.io/jq/manual/#example6
-- >>> query (JQCommaFilter (JQFieldFilter JQDotFilter "foo" False) (JQFieldFilter JQDotFilter "bar" False)) (fromJust $ JSON.decode "{\"foo\": 42, \"bar\": \"something else\", \"baz\": true}")
-- Right [Number 42.0,String "something else"]
-- >>> query (JQCommaFilter (JQFieldFilter JQDotFilter "user" False) (JQAllFilter (JQFieldFilter JQDotFilter "projects" False) False)) (fromJust $ JSON.decode "{\"user\":\"stedolan\", \"projects\": [\"jq\", \"wikiflow\"]}")
-- Right [String "stedolan",String "jq",String "wikiflow"]
-- >>> query (JQCommaFilter (JQFieldFilter JQDotFilter "user" False) (JQAllFilter (JQFieldFilter JQDotFilter "projects" False) False)) (fromJust $ JSON.decode "{\"user\":\"stedolan\", \"projects\": [\"jq\", \"wikiflow\"]}")
-- Right [String "stedolan",String "jq",String "wikiflow"]
--
-- http://stedolan.github.io/jq/manual/#example7
-- >>> query (JQBarFilter (JQAllFilter JQDotFilter False) (JQFieldFilter JQDotFilter "name" False)) (fromJust $ JSON.decode "[{\"name\":\"JSON\", \"good\":true}, {\"name\":\"XML\", \"good\":false}]")
-- Right [String "JSON",String "XML"]
--
-- http://stedolan.github.io/jq/manual/#example8
-- >>> query (JQArray [JQFieldFilter JQDotFilter "user" False, JQAllFilter (JQFieldFilter JQDotFilter "projects" False) False]) (fromJust $ JSON.decode "{\"user\":\"stedolan\", \"projects\": [\"jq\", \"wikiflow\"]}")
-- Right [Array (fromList [String "stedolan",String "jq",String "wikiflow"])]
--
query :: (Applicative m, MonadReader JSON.Value m) => JQFilter -> m (Either JQError [JSON.Value])
query JQDotFilter = Right . (:[]) <$> ask
query (JQFieldFilter jq key p)= (>>= \vs -> concat <$> mapM go vs) <$> query jq where
  go :: JSON.Value -> Either JQError [JSON.Value]
  go value@(JSON.Object _) = return [maybe JSON.Null id . JSON.parseMaybe (JSON.parseJSON >=> (JSON..: key)) $ value]
  go _ = if p then return [] else Left JQE
query (JQArrayFilter jq expr p) = (>>= \vs -> concat <$> mapM (go expr) vs) <$> query jq where
  go :: JQArrayFilterExpr -> JSON.Value -> Either JQError [JSON.Value]
  go (JQArrayFilterKey key) value@(JSON.Object _) = return [maybe JSON.Null id . JSON.parseMaybe (JSON.parseJSON >=> (JSON..: key)) $ value]
  go (JQArrayFilterAt ix) (JSON.Array array) = return [maybe JSON.Null id $ array V.!? ix]
  go (JQArrayFilterFromTo Nothing Nothing) _ = Left JQE
  go (JQArrayFilterFromTo mfrom mto) (JSON.Array array) = return . drop from . take to $ V.toList array where
      len = V.length array
      from = maybe 0 (\x -> x + if x < 0 then len else 0) mfrom
      to = maybe len (\x -> x + if x < 0 then len else 0) mto
  go (JQArrayFilterFromTo mfrom mto) (JSON.String str) = return [ JSON.toJSON $ T.drop from $ T.take to str ] where
      len = T.length str
      from = maybe 0 (\x -> x + if x < 0 then len else 0) mfrom
      to = maybe len (\x -> x + if x < 0 then len else 0) mto
  go _ _ = if p then return [] else Left JQE
query (JQAllFilter jq p) = (>>= \vs -> concat <$> mapM go vs) <$> query jq where
  go :: JSON.Value -> Either JQError [JSON.Value]
  go (JSON.Object object) = return $ HM.elems object
  go (JSON.Array array) = return $ V.toList array
  go _ = if p then return [] else Left JQE
query (JQCommaFilter a b) = (\x y -> (++) <$> x <*> y) <$> query a <*> query b
query (JQBarFilter a b)   = do
  evs <- query a
  case evs of
    Left e -> return $ Left e
    Right values -> foldl1' (\x y -> (++) <$> x <*> y) <$> mapM (\value -> local (const value) $ query b) values
query (JQArray jqs) = either Left (Right . (:[]) . JSON.toJSON) . foldl1' (\x y -> (++) <$> x <*> y) <$> mapM query jqs
query (JQObject elems) = undefined -- do objs <- mapM_ query' elems


query' (JQObjectShortcutElem object) = query' (JQObjectElem (Left object) (Right $ JQFieldFilter JQDotFilter object False))
query' (JQObjectElem ekeyjq evaluejq) = do
  ekey <- either (return . Right . (:[]) . JSON.toJSON) query ekeyjq
  case ekey of
    Right [JSON.String key] -> do
      evalue <- either (return . Right . (:[]) . JSON.toJSON) query evaluejq
      case evalue of
        Right value -> return $ Right [JSON.object [ key JSON..= value ] ]
        _           -> return $ Left JQE
    _                       -> return $ Left JQE